Agent: add support for remote server

comps/agent/src/README.md

@@ -82,7 +82,7 @@ for line in resp.iter_lines(decode_unicode=True):
 
 **Note**:
 
-1. Currently only `reract_llama` agent is enabled for assistants APIs.
+1. Currently only `react_llama` agent is enabled for assistants APIs.
 2. Not all keywords of OpenAI APIs are supported yet.
 
 ### 1.5 Agent memory
@@ -110,6 +110,32 @@ Examples of python code for multi-turn conversations using agent memory:
 
 To run the two examples above, first launch the agent microservice using [this docker compose yaml](../../../tests/agent/reactllama.yaml).
 
+### 1.6 Run LLMs from OpenAI
+
+To run any model from OpenAI, just specify the environment variable `OPENAI_API_KEY`:
+
+```bash
+export OPENAI_API_KEY=<openai-api-key>
+```
+
+These also need to be passed in to the `docker run` command, or included in a YAML file when running `docker compose`.
+
+### 1.7 Run LLMs with OpenAI-compatible APIs on Remote Servers
+
+To run the text generation portion using LLMs deployed on a remote server, specify the following environment variables:
+
+```bash
+export api_key=<openai-api-key>
+export model=<model-card>
+export LLM_ENDPOINT_URL=<inference-endpoint>
+```
+
+These also need to be passed in to the `docker run` command, or included in a YAML file when running `docker compose`.
+
+#### Notes
+
+- For `LLM_ENDPOINT_URL`, there is no need to include `v1`.
+
 ## 🚀2. Start Agent Microservice
 
 ### 2.1 Build docker image for agent microservice

comps/agent/src/integrations/config.py

@@ -17,6 +17,9 @@
 if os.environ.get("llm_endpoint_url") is not None:
     env_config += ["--llm_endpoint_url", os.environ["llm_endpoint_url"]]
 
+if os.environ.get("api_key") is not None:
+    env_config += ["--api_key", os.environ["api_key"]]
+
 if os.environ.get("llm_engine") is not None:
     env_config += ["--llm_engine", os.environ["llm_engine"]]
 

comps/agent/src/integrations/utils.py

@@ -66,7 +66,17 @@ def setup_chat_model(args):
             **params,
         )
     elif args.llm_engine == "openai":
-        llm = ChatOpenAI(model_name=args.model, request_timeout=args.timeout, **params)
+        if args.api_key:
+            openai_endpoint = f"{args.llm_endpoint_url}/v1"
+            llm = ChatOpenAI(
+                openai_api_key=args.api_key,
+                openai_api_base=openai_endpoint,
+                model_name=args.model,
+                request_timeout=args.timeout,
+                **params,
+            )
+        else:
+            llm = ChatOpenAI(model_name=args.model, request_timeout=args.timeout, **params)
     else:
         raise ValueError("llm_engine must be vllm, tgi or openai")
     return llm
@@ -162,6 +172,7 @@ def get_args():
     parser.add_argument("--model", type=str, default="meta-llama/Meta-Llama-3-8B-Instruct")
     parser.add_argument("--llm_engine", type=str, default="tgi")
     parser.add_argument("--llm_endpoint_url", type=str, default="http://localhost:8080")
+    parser.add_argument("--api_key", type=str, default=None, help="API key to access remote server")
     parser.add_argument("--max_new_tokens", type=int, default=1024)
     parser.add_argument("--top_k", type=int, default=10)
     parser.add_argument("--top_p", type=float, default=0.95)