Update src/agents/rag_agent.py

src/agents/rag_agent.py

@@ -62,21 +62,16 @@ tools = [search_legal_documents, search_youtube_transcripts]
 llm = GeminiWrapper()
 llm_with_tools = llm.bind_tools(tools)
 
-# SYSTEM MESSAGE
-#sys_msg = SystemMessage(content="""
-#You are a helpful assistant specialized in answering user questions related to Moroccan Parliament YouTube videos and legal documents.
-#Your response must be strictly in the same language as the user’s query.
-#Provide accurate answers and include relevant sources (YouTube video links or PDF document links) in your response.
-#""")
-
 # LLM NODE
 def assistant(state: MessagesState):
      # Utiliser l'état d'invocation global
-    invocation_type = invocation_state.invocation_type
+    #invocation_type = invocation_state.invocation_type
+
+    user_msg = state["messages"][-1].content
+    logging.info(f"User input: {user_msg}")
 
-    print(invocation_type)
     # Créer le SystemMessage selon le type d'invocation
-    if invocation_type == "chatbot":
+    if invocation_state.invocation_type == 'chatbot':
         sys_msg = SystemMessage(content="""
         You are a helpful assistant specialized in answering user questions related to Moroccan Parliament YouTube videos and legal documents.
         Your response must be strictly in the same language as the user’s query.
@@ -88,15 +83,28 @@ def assistant(state: MessagesState):
         Ensure your responses are concise and formatted appropriately for API output. Your response should be maximum 100 words.
         Your response shoud be in json format like that : {"text response" : "", "sources" : ""}
         """)
-    user_msg = state["messages"][-1].content
-    logging.info(f"User input: {user_msg}")
+    
     try:
+        print(state['messages'])
         result = llm_with_tools.invoke([sys_msg] + state["messages"])
-        logging.info(f"🤖 Model Output: {result}")
+        logging.info(f"🤖 Model Output kkkk: {result}")
         return {"messages": [result]}
     except Exception as e:
         logging.error(f"LLM invocation failed: {e}")
         return {"messages": [SystemMessage(content="Une erreur est survenue avec le modèle.")]}
+    
+# Tool to set the invocation type based on the user's message
+def set_invocation_type(state: MessagesState):
+    user_msg = state["messages"][-1].content.strip()  # Normalize the input
+    logging.info(f"User input for type update: {user_msg}")
+
+    # Set the invocation type based on the presence of 'APICALL'
+    if 'apicall' in user_msg:
+        invocation_state.invocation_type = 'API'
+    else:
+        invocation_state.invocation_type = 'chatbot'
+
+    logging.info(f"Invocation type set to: {invocation_state.invocation_type}")
 
 class InvocationState:
     def __init__(self):
@@ -108,10 +116,14 @@ invocation_state = InvocationState()
 # GRAPH SETUP
 builder = StateGraph(MessagesState)
 
+# Add the node to set the invocation type
+builder.add_node("set_invocation_type", set_invocation_type)
 builder.add_node("llm_assistant", assistant)
 builder.add_node("tools", ToolNode(tools))
 
-builder.add_edge(START, "llm_assistant")
+# Define the edges of the graph
+builder.add_edge(START, "set_invocation_type")
+builder.add_edge("set_invocation_type", "llm_assistant")
 builder.add_conditional_edges("llm_assistant", tools_condition)
 builder.add_edge("tools", "llm_assistant")