feat: generating final answer if specified

basic_agent.py

@@ -40,6 +40,31 @@ def print_conversation(messages):
         console.print(panel)
 
 
+def generate_final_answer(qa: dict[str, str]) -> str:
+    """Invokes gpt-4o-mini to extract generate a final answer based on the content query, response, and metadata"""
+
+    final_answer_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+    system_prompt = (
+        "You will receive a JSON string containing a user's query, a response, and metadata. "
+        "Extract and return only the final answer to the query as a plain string. "
+        "Do not return anything else. "
+        "Avoid any labels, prefixes, or explanation. "
+        "Return only the exact value that satisfies the query, suitable for string comparison."
+        "If the query is not answerable due to a missing file in the input and is reflected in the response, answer with 'File not found'. "
+    )
+
+    system_message = SystemMessage(content=system_prompt)
+    messages = [
+        system_message,
+        HumanMessage(content=f'Generate the final answer for the following query:\n\n{json.dumps(qa)}')
+    ]
+
+    response = final_answer_llm.invoke(messages)
+
+    return response.content
+    
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], add_messages]
 
@@ -104,7 +129,7 @@ class BasicOpenAIAgentWorkflow:
         self.graph = graph.compile()
 
 
-    def chat(self, query, verbose=2):
+    def chat(self, query, verbose=2, only_final_answer=False):
         """Simple agent call"""
         if isinstance(query, dict):
             query = query["messages"]
@@ -121,16 +146,25 @@ class BasicOpenAIAgentWorkflow:
         if not self.history_messages:
             raise ValueError("Converted message history is empty. Something went wrong.")
 
-        response = self.graph.invoke({'messages': self.history_messages})  # invoke with all the history
+        response = self.graph.invoke({'messages': self.history_messages})  # invoke with all the history to keep context (dummy mem)
         response = response['messages'][-1].content
+        
+        if only_final_answer:
+            final_answer_content = {
+                'query': query,
+                'response': response,
+                'metadata': {}
+            }
+            response = generate_final_answer(final_answer_content)
+
         assistant_message = {'role': 'assistant', 'content': response}
         self.history.append(assistant_message)
-        
+
         if verbose==2:
             print_conversation(self.history)
         elif verbose==1:
             print_conversation([response])
-            
+
         return response
 
 
@@ -141,12 +175,12 @@ class BasicOpenAIAgentWorkflow:
         return {'messages': self.history_messages}
 
 
-    def chat_batch(self, queries=None):
+    def chat_batch(self, queries=None, only_final_answer=False):
         """Send several simple agent calls to the llm using the compiled graph"""
         if queries is None:
             queries = []
         for i, query in tqdm(enumerate(queries, start=1)):
             if i == len(queries):
-                self.chat(query, verbose=2)
+                self.chat(query, verbose=2, only_final_answer=only_final_answer)
             else:
-                self.chat(query, verbose=0)
+                self.chat(query, verbose=0, only_final_answer=only_final_answer)