Upload certificate_agent.py

certificate_agent.py

@@ -0,0 +1,106 @@
+from langgraph.graph import StateGraph, END
+from typing import TypedDict, List, Literal
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+from visit_web_pages_tool import visit_webpage
+import wikipedia
+import json
+from prompt import SYSTEM_PROMPT_MANAGER, SYSTEM_PROMPT_CLEANER
+from langchain_community.tools import DuckDuckGoSearchResults
+
+class GraphState(TypedDict):
+    history: List
+
+llm = HuggingFaceEndpoint(
+    repo_id="Qwen/Qwen2.5-72B-Instruct",
+    task='text-generation',
+    max_new_tokens=4096
+)
+manager_agent = ChatHuggingFace(llm=llm)
+#cleaner_agent = ChatHuggingFace(llm=llm)
+
+def llm_call(state: GraphState) -> GraphState:
+    """
+        Node used to generate the basic LLM calls from the manager agent.
+    """
+    print(state['history'])
+    answer_llm = manager_agent.invoke(state['history'])
+    state['history'].append(answer_llm)
+    return state
+
+def tool_call(state: GraphState) -> GraphState:
+    """
+        Node used to perform tool call. For the moment, the only tool available is web_research.
+    """
+    #Première étape, convertir le dernier call en json
+    json_last_answer = json.loads(state['history'][-1].content)
+    if (json_last_answer['action'] == 'web_search'):
+        result_search = wikipedia.search(json_last_answer['query'])
+        markdown_website = visit_webpage(wikipedia.page(result_search[0]).url)
+        state['history'].append(
+            {'role': 'tool', 'content': markdown_website, 'tool_call_id': 'blablabla'}
+        )
+        return state
+    else:
+        state['history'].append(
+            {'role': 'tool', 'content': 'Invalid tool call', 'tool_call_id': 'blablabla'}
+        )
+        return state
+
+def force_final_answer(state: GraphState) -> GraphState:
+    state['history'].append(
+        {'role': 'human', 'content': 'Now provide the final answer based on the intermediate answer'}
+    )
+    return state
+
+def router_edge_tool(state: GraphState) -> Literal["llm_call", "tool_call", "end"]:
+    json_last_answer = json.loads(state['history'][-1].content)
+    if (json_last_answer['action'] == "intermediate_answer"):
+        return "force_final_answer"
+    elif (json_last_answer['action'] == 'web_search'):
+        return "tool_call"
+    else:
+        return "end"
+
+
+my_graph_build = StateGraph(GraphState)
+my_graph_build.add_node("llm_call", llm_call)
+my_graph_build.add_node("tool_call", tool_call)
+my_graph_build.add_node("force_final_answer", force_final_answer)
+my_graph_build.add_conditional_edges("llm_call", router_edge_tool, {"force_final_answer": "force_final_answer", "tool_call": "tool_call", "end": END})
+my_graph_build.add_edge("tool_call", "llm_call")
+my_graph_build.add_edge("force_final_answer", "llm_call")
+my_graph_build.set_entry_point("llm_call")
+my_graph = my_graph_build.compile()
+
+init_state = GraphState(history=[
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': 'How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?'}
+    ])
+
+init_state_2 = GraphState(history=[
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': 'Who did the actor who played Ray in the Polish-language version of Everybody Loves Raymond play in Magda M.? Give only the first name.'}
+    ])
+
+init_state_3 = GraphState(history=[
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': 'What is the first name of the only Malko Competition recipient from the 20th Century (after 1977) whose nationality on record is a country that no longer exists?'}
+    ])
+
+init_state_4 = GraphState(history=[
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': '.rewsna eht sa \"tfel\" drow eht fo etisoppo eht etirw ,ecnetnes siht dnatsrednu uoy fI'}
+    ])
+
+init_state_5 = GraphState(history=[
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': "What country had the least number of athletes at the 1928 Summer Olympics? If there's a tie for a number of athletes, return the first in alphabetical order. Give the IOC country code as your answer."}
+    ])
+"""
+print(manager_agent.invoke([
+    {'role': 'system', 'content': SYSTEM_PROMPT_MANAGER},
+    {'role': 'human', 'content': 'What is the first name of the only Malko Competition recipient from the 20th Century (after 1977) whose nationality on record is a country that no longer exists?'}
+    ]))
+"""
+#print(my_graph.invoke(init_state_2))
+#print(my_graph.invoke(init_state_5))