Various small fixes to the serve application (#27)

- Make sure the right API is used
- Make sure the DB connection string is set correctly

Author: pcmoritz

app/query.py

@@ -6,8 +6,7 @@
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from pgvector.psycopg import register_vector
 
-from app.config import DB_CONNECTION_STRING
-
+import app.config
 
 def generate_response(
     llm,
@@ -57,7 +56,7 @@ def __init__(
         self.assistant_content = assistant_content
 
         # VectorDB connection
-        self.conn = psycopg.connect(DB_CONNECTION_STRING)
+        self.conn = psycopg.connect(app.config.DB_CONNECTION_STRING)
         register_vector(self.conn)
 
     def get_response(self, query):

app/serve.py

@@ -1,9 +1,10 @@
 # You can run the whole script locally with
-# serve run serve:deployment
+# serve run app.serve:deployment
 
 import json
 import os
 
+import openai
 import ray
 import requests
 from fastapi import FastAPI
@@ -17,18 +18,10 @@
 
 
 def get_secret(secret_name):
-    aws_secret_id = os.environ.get("RAY_ASSISTANT_AWS_SECRET_ID")
-    if aws_secret_id:
-        import boto3
-        client = boto3.client(
-            "secretsmanager", region_name=os.environ["RAY_ASSISTANT_AWS_REGION"]
-        )
-        response = client.get_secret_value(SecretId=aws_secret_id)
-        return json.loads(response["SecretString"])[secret_name]
-    else:
-        raise NotImplemented(
-            "Currently only AWS is supported "
-            "and you need to set RAY_ASSISTANT_AWS_SECRET_ID")
+    import boto3
+    client = boto3.client("secretsmanager", region_name="us-west-2")
+    response = client.get_secret_value(SecretId="ray-assistant")
+    return json.loads(response["SecretString"])[secret_name]
 
 
 application = FastAPI()
@@ -69,6 +62,8 @@ class Answer(BaseModel):
 class RayAssistantDeployment:
     def __init__(self):
         app.config.DB_CONNECTION_STRING = get_secret("DB_CONNECTION_STRING")
+        openai.api_key = get_secret("OPENAI_API_KEY")
+        openai.api_base = "https://api.endpoints.anyscale.com/v1"
         self.agent = query.QueryAgent(
             llm="meta-llama/Llama-2-70b-chat-hf",
             max_context_length=4096,

app/service.yaml

@@ -3,8 +3,4 @@ cluster_env: ray-assistant:2
 ray_serve_config:
   import_path: app.serve:deployment
   runtime_env:
-    working_dir: "https://github.com/ray-project/llm-applications/archive/refs/tags/v0.0.5.zip"
-    env_vars: {
-      RAY_ASSISTANT_AWS_SECRET_ID: "ray-assistant",
-      RAY_ASSISTANT_AWS_REGION: "us-west-2"
-    }
+    working_dir: "https://github.com/ray-project/llm-applications/archive/refs/tags/v0.0.6.zip"