Spaces:
Sleeping
Sleeping
| import pinecone | |
| import openai | |
| from openai import OpenAI | |
| from Answering_Agent import Answering_Agent | |
| from Obnoxious_Agent import Obnoxious_Agent | |
| from Query_Agent import Query_Agent | |
| from Relevant_Documents_Agent import Relevant_Documents_Agent | |
| class Head_Agent: | |
| def __init__(self, openai_key, pinecone_key, pinecone_index_name) -> None: | |
| # TODO: Initialize the Head_Agent | |
| # Store API keys and Pinecone index name | |
| openai_key = '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' | |
| self.openai_client = OpenAI(api_key=openai_key) | |
| self.openai_key = openai_key | |
| self.pinecone_key = pinecone_key | |
| self.pinecone_index_name = pinecone_index_name | |
| # Initialize Pinecone | |
| # Placeholder for sub-agents, to be initialized in setup_sub_agents | |
| self.obnoxious_agent = None | |
| self.query_agent = None | |
| self.answering_agent = None | |
| self.relevant_documents_agent = None | |
| # set up sub agents | |
| self.setup_sub_agents() | |
| def get_completion(self, prompt, model="gpt-4"): | |
| message = {"role": "user", "content": prompt} | |
| response = self.openai_client.chat.completions.create( | |
| model=model, | |
| messages=[message] | |
| ) | |
| return response.choices[0].message.content | |
| def setup_sub_agents(self): | |
| # TODO: Setup the sub-agents | |
| # Setup the Obnoxious_Agent | |
| self.obnoxious_agent = Obnoxious_Agent() # Assuming no external API client required | |
| # Set up the Query_Agent with Pinecone index | |
| self.query_agent = Query_Agent(pinecone_index=self.pinecone_index_name, openai_client=openai, | |
| embeddings="text-embedding-ada-002") | |
| # Set up the Answering_Agent | |
| self.answering_agent = Answering_Agent(openai_api_key=self.openai_key) | |
| openai_client = OpenAI(api_key=self.openai_key) | |
| self.relevant_documents_agent = Relevant_Documents_Agent(openai_client) | |
| def process_query(self, user_input, conversation_history): | |
| # Check if the query is obnoxious | |
| # conversation_history = get_conversation() | |
| if self.obnoxious_agent.check_query(user_input) == "Yes": | |
| return "Sorry, I cannot respond to this query." | |
| greetings = [ | |
| "hi", "hello", "how are you", "hey", "good morning", | |
| "good afternoon", "good evening", "greetings", "what's up", | |
| "howdy", "hi there", "hello there", "hey there" | |
| ] | |
| if user_input.lower() in greetings: | |
| return "How can I help you today?" | |
| # Retrieve relevant documents for the query (simplified example) | |
| # In a real scenario, you might need to preprocess the query for embeddings | |
| relevant_docs = self.query_agent.query_vector_store(user_input, k=5) | |
| print("relevant_docs .length : ") | |
| print(len(relevant_docs)) | |
| if len(relevant_docs) == 0: | |
| none_relevant_docs = "No relevant documents found in the documents.Please ask a relevant question to the book on Machine Learning." | |
| return none_relevant_docs | |
| # choose mode | |
| prompt_choose_mode = "Infer whether the user wants the response to be chatty or concise? If the response needs to be chatty you must answer chatty, otherwise answer concise and the answer must contain only the one word. Here is user_input: "+ user_input | |
| prompt_choose_mode = self.get_completion(prompt=prompt_choose_mode, model="gpt-4") | |
| if prompt_choose_mode == "chatty": | |
| response = self.answering_agent.generate_response(user_input, relevant_docs, conv_history=conversation_history, | |
| mode="chatty") | |
| elif prompt_choose_mode =="concise": | |
| response = self.answering_agent.generate_response(user_input, relevant_docs, | |
| conv_history=conversation_history, | |
| mode="concise") | |
| else: | |
| # Generate a response based on the query and relevant documents | |
| response = self.answering_agent.generate_response(user_input, relevant_docs, conv_history=conversation_history, | |
| mode="chatty") | |
| print("prompt_choose_mode: "+ prompt_choose_mode ) | |
| return response | |