Added simple backend agent api

codingbandit · codingbandit · commit 28767b33dad8 · 2023-12-27T21:35:58.000-05:00
Includes in-memory conversation history
diff --git a/Labs/app/backend/api_models/ai_request.py b/Labs/app/backend/api_models/ai_request.py
@@ -0,0 +1,13 @@
+"""
+AIRequest model
+"""
+from pydantic import BaseModel
+
+class AIRequest(BaseModel):
+    """
+    AIRequest model encapsulates the session_id
+    and incoming user prompt for the AI agent
+    to respond to.
+    """
+    session_id: str
+    prompt: str
diff --git a/Labs/app/backend/cosmic_works/cosmic_works_ai_agent.py b/Labs/app/backend/cosmic_works/cosmic_works_ai_agent.py
@@ -0,0 +1,178 @@
+"""
+The CosmicWorksAIAgent class encapsulates a LangChain 
+agent that can be used to answer questions about Cosmic Works
+products, customers, and sales.
+"""
+import os
+import json
+from typing import List
+import pymongo
+from dotenv import load_dotenv
+from langchain.chat_models import AzureChatOpenAI
+from langchain.embeddings import AzureOpenAIEmbeddings
+from langchain.vectorstores.azure_cosmos_db import AzureCosmosDBVectorSearch
+from langchain.schema.document import Document
+from langchain.agents import Tool
+from langchain.agents.agent_toolkits import create_conversational_retrieval_agent
+from langchain.tools import StructuredTool
+from langchain_core.messages import SystemMessage
+
+load_dotenv("../../.env")
+DB_CONNECTION_STRING = os.environ.get("DB_CONNECTION_STRING")
+db = pymongo.MongoClient(DB_CONNECTION_STRING).cosmic_works
+
+class CosmicWorksAIAgent:
+    """
+    The CosmicWorksAIAgent class creates Cosmo, an AI agent
+    that can be used to answer questions about Cosmic Works
+    products, customers, and sales.
+    """
+    def __init__(self, session_id: str):
+        llm = AzureChatOpenAI(
+            temperature = 0,
+            openai_api_version = "2023-09-01-preview",
+            azure_endpoint = os.environ.get("AOAI_ENDPOINT"),
+            openai_api_key = os.environ.get("AOAI_KEY"),
+            azure_deployment = "completions"
+        )
+        self.embedding_model = AzureOpenAIEmbeddings(
+            openai_api_version = "2023-09-01-preview",
+            azure_endpoint = os.environ.get("AOAI_ENDPOINT"),
+            openai_api_key = os.environ.get("AOAI_KEY"),
+            azure_deployment = "embeddings",
+            chunk_size=10
+        )
+        system_message = SystemMessage(
+            content = """
+                You are a helpful, fun and friendly sales assistant for Cosmic Works, 
+                a bicycle and bicycle accessories store.
+
+                Your name is Cosmo.
+
+                You are designed to answer questions about the products that Cosmic Works sells, 
+                the customers that buy them, and the sales orders that are placed by customers.
+
+                If you don't know the answer to a question, respond with "I don't know."
+            """
+        )
+        self.agent_executor = create_conversational_retrieval_agent(
+                llm,
+                self.__create_agent_tools(),
+                system_message = system_message,
+                memory_key=session_id,
+                verbose=True
+        )
+
+    def run(self, prompt: str) -> str:
+        """
+        Run the AI agent.
+        """
+        result = self.agent_executor({"input": prompt})
+        return result["output"]
+
+    def __create_cosmic_works_vector_store_retriever(
+            self,
+            collection_name: str,
+            top_k: int = 3
+        ):
+        """
+        Returns a vector store retriever for the given collection.
+        """
+        vector_store =  AzureCosmosDBVectorSearch.from_connection_string(
+            connection_string = DB_CONNECTION_STRING,
+            namespace = f"cosmic_works.{collection_name}",
+            embedding = self.embedding_model,
+            index_name = "VectorSearchIndex",
+            embedding_key = "contentVector",
+            text_key = "_id"
+        )
+        return vector_store.as_retriever(search_kwargs={"k": top_k})
+
+    def __create_agent_tools(self) -> List[Tool]:
+        """
+        Returns a list of agent tools.
+        """
+        products_retriever = self.__create_cosmic_works_vector_store_retriever("products")
+        customers_retriever = self.__create_cosmic_works_vector_store_retriever("customers")
+        sales_retriever = self.__create_cosmic_works_vector_store_retriever("sales")
+
+        # create a chain on the retriever to format the documents as JSON
+        products_retriever_chain = products_retriever | format_docs
+        customers_retriever_chain = customers_retriever | format_docs
+        sales_retriever_chain = sales_retriever | format_docs
+
+        tools = [
+            Tool(
+                name = "vector_search_products",
+                func = products_retriever_chain.invoke,
+                description = """
+                    Searches Cosmic Works product information for similar products based 
+                    on the question. Returns the product information in JSON format.
+                    """
+            ),
+            Tool(
+                name = "vector_search_customers",
+                func = customers_retriever_chain.invoke,
+                description = """
+                    Searches Cosmic Works customer information and retrieves similar 
+                    customers based on the question. Returns the customer information 
+                    in JSON format.
+                    """
+            ),
+            Tool(
+                name = "vector_search_sales",
+                func = sales_retriever_chain.invoke,
+                description = """
+                    Searches Cosmic Works customer sales information and retrieves sales order 
+                    details based on the question. Returns the sales order information in JSON format.
+                    """
+            ),
+            StructuredTool.from_function(get_product_by_id),
+            StructuredTool.from_function(get_product_by_sku),
+            StructuredTool.from_function(get_sales_by_id)
+        ]
+        return tools
+
+def format_docs(docs:List[Document]) -> str:
+    """
+    Prepares the product list for the system prompt.
+    """
+    str_docs = []
+    for doc in docs:
+        # Build the product document without the contentVector
+        doc_dict = {"_id": doc.page_content}
+        doc_dict.update(doc.metadata)
+        if "contentVector" in doc_dict:
+            del doc_dict["contentVector"]
+        str_docs.append(json.dumps(doc_dict, default=str))
+    # Return a single string containing each product JSON representation
+    # separated by two newlines
+    return "\n\n".join(str_docs)
+
+def get_product_by_id(product_id: str) -> str:
+    """
+    Retrieves a product by its ID.    
+    """
+    doc = db.products.find_one({"_id": product_id})
+    if "contentVector" in doc:
+        del doc["contentVector"]
+    return json.dumps(doc)
+
+def get_product_by_sku(sku: str) -> str:
+    """
+    Retrieves a product by its sku.
+    """
+    doc = db.products.find_one({"sku": sku})
+    if "contentVector" in doc:
+        del doc["contentVector"]
+    return json.dumps(doc, default=str)
+
+def get_sales_by_id(sales_id: str) -> str:
+    """
+    Retrieves a sales order by its ID.
+    """
+    doc = db.sales.find_one({"_id": sales_id})
+    if "contentVector" in doc:
+        del doc["contentVector"]
+    return json.dumps(doc, default=str)
+  
diff --git a/Labs/app/backend/main.py b/Labs/app/backend/main.py
@@ -0,0 +1,30 @@
+"""
+API entrypoint for backend API.
+"""
+import uvicorn
+from fastapi import FastAPI
+from api_models.ai_request import AIRequest
+from cosmic_works.cosmic_works_ai_agent import CosmicWorksAIAgent
+
+app = FastAPI()
+# Agent pool keyed by session_id to retain memories/history in-memory.
+agent_pool = {}
+
+@app.get("/")
+def root():
+    """
+    Health probe endpoint.
+    """
+    return {"status": "ready"}
+
+@app.post("/ai")
+def run_cosmic_works_ai_agent(request: AIRequest):
+    """
+    Run the Cosmic Works AI agent.
+    """
+    if request.session_id not in agent_pool:
+        agent_pool[request.session_id] = CosmicWorksAIAgent(request.session_id)
+    return { "message": agent_pool[request.session_id].run(request.prompt) }
+
+if __name__ == "__main__":
+    uvicorn.run("main:app", host="0.0.0.0", port=4242, reload=True)
diff --git a/Labs/requirements.txt b/Labs/requirements.txt
@@ -5,4 +5,6 @@ pydantic==2.5.2
 openai==1.6.0
 tenacity==8.2.3
 langchain==0.0.352
-tiktoken==0.5.2
+tiktoken==0.5.2
+fastapi==0.108.0
+uvicorn==0.25.0
