Update agent_v2.py

agent_v2.py

@@ -13,6 +13,8 @@ import pandas as pd
 import uuid
 import numpy as np
 from code_interpreter import CodeInterpreter
+import re
+from langchain_core.messages import AIMessage
 
 interpreter_instance = CodeInterpreter()
 
@@ -747,6 +749,17 @@ def build_graph(provider: str = "huggingface"):
         """Assistant node"""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
+    # Define extractor node    
+    def extractor(state: MessagesState):
+        last_message = state["messages"][-1].content
+        match = re.search(r"FINAL ANSWER:\s*(.*)", last_message, re.IGNORECASE)
+        if match:
+            answer = match.group(1).strip().split("\n")[0].rstrip('.')
+            return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
+        else:
+            return {"messages": [AIMessage(content="FINAL ANSWER: No valid answer found")]}
+
+
     """
     def retriever(state: MessagesState):
         "Retriever node"
@@ -762,6 +775,7 @@ def build_graph(provider: str = "huggingface"):
             return {"messages": [sys_msg] + state["messages"]}
     """
 
+    """
     builder = StateGraph(MessagesState)
     #builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
@@ -774,6 +788,18 @@ def build_graph(provider: str = "huggingface"):
         tools_condition,
     )
     builder.add_edge("tools", "assistant")
+    """
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_node("extractor", extractor)
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+    builder.add_edge("assistant", "extractor")
+    builder.set_finish_point("extractor")
+
 
     # Compile graph
     return builder.compile()