1+ """
2+ The CosmicWorksAIAgent class encapsulates a LangChain
3+ agent that can be used to answer questions about Cosmic Works
4+ products, customers, and sales.
5+ """
6+ import os
7+ import json
8+ from typing import List
9+ import pymongo
10+ from dotenv import load_dotenv
11+ from langchain .chat_models import AzureChatOpenAI
12+ from langchain .embeddings import AzureOpenAIEmbeddings
13+ from langchain .vectorstores .azure_cosmos_db import AzureCosmosDBVectorSearch
14+ from langchain .schema .document import Document
15+ from langchain .agents import Tool
16+ from langchain .agents .agent_toolkits import create_conversational_retrieval_agent
17+ from langchain .tools import StructuredTool
18+ from langchain_core .messages import SystemMessage
19+
20+ load_dotenv (".env" )
21+ DB_CONNECTION_STRING = os .environ .get ("DB_CONNECTION_STRING" )
22+ AOAI_ENDPOINT = os .environ .get ("AOAI_ENDPOINT" )
23+ AOAI_KEY = os .environ .get ("AOAI_KEY" )
24+ AOAI_API_VERSION = "2023-09-01-preview"
25+ COMPLETIONS_DEPLOYMENT = "completions"
26+ EMBEDDINGS_DEPLOYMENT = "embeddings"
27+ db = pymongo .MongoClient (DB_CONNECTION_STRING ).cosmic_works
28+
29+ class CosmicWorksAIAgent :
30+ """
31+ The CosmicWorksAIAgent class creates Cosmo, an AI agent
32+ that can be used to answer questions about Cosmic Works
33+ products, customers, and sales.
34+ """
35+ def __init__ (self , session_id : str ):
36+ llm = AzureChatOpenAI (
37+ temperature = 0 ,
38+ openai_api_version = AOAI_API_VERSION ,
39+ azure_endpoint = AOAI_ENDPOINT ,
40+ openai_api_key = AOAI_KEY ,
41+ azure_deployment = COMPLETIONS_DEPLOYMENT
42+ )
43+ self .embedding_model = AzureOpenAIEmbeddings (
44+ openai_api_version = AOAI_API_VERSION ,
45+ azure_endpoint = AOAI_ENDPOINT ,
46+ openai_api_key = AOAI_KEY ,
47+ azure_deployment = EMBEDDINGS_DEPLOYMENT ,
48+ chunk_size = 10
49+ )
50+ system_message = SystemMessage (
51+ content = """
52+ You are a helpful, fun and friendly sales assistant for Cosmic Works,
53+ a bicycle and bicycle accessories store.
54+
55+ Your name is Cosmo.
56+
57+ You are designed to answer questions about the products that Cosmic Works sells,
58+ the customers that buy them, and the sales orders that are placed by customers.
59+
60+ If you don't know the answer to a question, respond with "I don't know."
61+
62+ Only answer questions related to Cosmic Works products, customers, and sales orders.
63+
64+ If a question is not related to Cosmic Works products, customers, or sales orders,
65+ respond with "I only answer questions about Cosmic Works"
66+ """
67+ )
68+ self .agent_executor = create_conversational_retrieval_agent (
69+ llm ,
70+ self .__create_agent_tools (),
71+ system_message = system_message ,
72+ memory_key = session_id ,
73+ verbose = True
74+ )
75+
76+ def run (self , prompt : str ) -> str :
77+ """
78+ Run the AI agent.
79+ """
80+ result = self .agent_executor ({"input" : prompt })
81+ return result ["output" ]
82+
83+ def __create_cosmic_works_vector_store_retriever (
84+ self ,
85+ collection_name : str ,
86+ top_k : int = 3
87+ ):
88+ """
89+ Returns a vector store retriever for the given collection.
90+ """
91+ vector_store = AzureCosmosDBVectorSearch .from_connection_string (
92+ connection_string = DB_CONNECTION_STRING ,
93+ namespace = f"cosmic_works.{ collection_name } ,
94+ embedding = self .embedding_model ,
95+ index_name = "VectorSearchIndex" ,
96+ embedding_key = "contentVector" ,
97+ text_key = "_id"
98+ )
99+ return vector_store .as_retriever (search_kwargs = {"k" : top_k })
100+
101+ def __create_agent_tools (self ) -> List [Tool ]:
102+ """
103+ Returns a list of agent tools.
104+ """
105+ products_retriever = self .__create_cosmic_works_vector_store_retriever ("products" )
106+ customers_retriever = self .__create_cosmic_works_vector_store_retriever ("customers" )
107+ sales_retriever = self .__create_cosmic_works_vector_store_retriever ("sales" )
108+
109+ # create a chain on the retriever to format the documents as JSON
110+ products_retriever_chain = products_retriever | format_docs
111+ customers_retriever_chain = customers_retriever | format_docs
112+ sales_retriever_chain = sales_retriever | format_docs
113+
114+ tools = [
115+ Tool (
116+ name = "vector_search_products" ,
117+ func = products_retriever_chain .invoke ,
118+ description = """
119+ Searches Cosmic Works product information for similar products based
120+ on the question. Returns the product information in JSON format.
121+ """
122+ ),
123+ Tool (
124+ name = "vector_search_customers" ,
125+ func = customers_retriever_chain .invoke ,
126+ description = """
127+ Searches Cosmic Works customer information and retrieves similar
128+ customers based on the question. Returns the customer information
129+ in JSON format.
130+ """
131+ ),
132+ Tool (
133+ name = "vector_search_sales" ,
134+ func = sales_retriever_chain .invoke ,
135+ description = """
136+ Searches Cosmic Works customer sales information and retrieves sales order
137+ details based on the question. Returns the sales order information in JSON format.
138+ """
139+ ),
140+ StructuredTool .from_function (get_product_by_id ),
141+ StructuredTool .from_function (get_product_by_sku ),
142+ StructuredTool .from_function (get_sales_by_id )
143+ ]
144+ return tools
145+
146+ def format_docs (docs :List [Document ]) -> str :
147+ """
148+ Prepares the product list for the system prompt.
149+ """
150+ str_docs = []
151+ for doc in docs :
152+ # Build the product document without the contentVector
153+ doc_dict = {"_id" : doc .page_content }
154+ doc_dict .update (doc .metadata )
155+ if "contentVector" in doc_dict :
156+ del doc_dict ["contentVector" ]
157+ str_docs .append (json .dumps (doc_dict , default = str ))
158+ # Return a single string containing each product JSON representation
159+ # separated by two newlines
160+ return "\n \n " .join (str_docs )
161+
162+ def get_product_by_id (product_id : str ) -> str :
163+ """
164+ Retrieves a product by its ID.
165+ """
166+ doc = db .products .find_one ({"_id" : product_id })
167+ if "contentVector" in doc :
168+ del doc ["contentVector" ]
169+ return json .dumps (doc )
170+
171+ def get_product_by_sku (sku : str ) -> str :
172+ """
173+ Retrieves a product by its sku.
174+ """
175+ doc = db .products .find_one ({"sku" : sku })
176+ if "contentVector" in doc :
177+ del doc ["contentVector" ]
178+ return json .dumps (doc , default = str )
179+
180+ def get_sales_by_id (sales_id : str ) -> str :
181+ """
182+ Retrieves a sales order by its ID.
183+ """
184+ doc = db .sales .find_one ({"_id" : sales_id })
185+ if "contentVector" in doc :
186+ del doc ["contentVector" ]
187+ return json .dumps (doc , default = str )
188+
0 commit comments