implement advanced LangGraph processing

workflow.py

@@ -1,12 +1,25 @@
-from typing import TypedDict, Annotated, Callable, Optional, Any
+from typing import List, Dict, TypedDict, Annotated, Callable, Optional, Any
+from huggingface_hub.inference._generated.types import question_answering
 from langgraph.graph import StateGraph, END
+from langchain_core.messages import HumanMessage
+#from langchain_community.chat_models import ChatHuggingFace
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from answering import gen_question_answer
 
+import requests
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
 
 class AgentState(TypedDict):
+    context: Dict[str, str]
     question: str
     answer: Annotated[str, lambda x, y: y]  # Overwrite with new value
     formatted_answer: Annotated[str, lambda x, y: y]
+    format_requirement: str    
+    uris: List[str]
+    reasoning: List[str]
 
 class GAIAAnsweringWorkflow:
     def __init__(
@@ -17,14 +30,40 @@ class GAIAAnsweringWorkflow:
         """
         Initialize the GAIA agent workflow
         
-        Args:
+        Args: 
             qa_function: Core question answering function (gen_question_answer)
             formatter: Answer formatting function (default: GAIA boxed format)
         """
         self.qa_function = gen_question_answer #qa_function or self.default_qa_function
         self.formatter = formatter or self.default_formatter
         self.workflow = self.build_workflow()
-        
+        # Initialize model with HF Inference API
+        llm_endpoint = HuggingFaceEndpoint(
+            model="deepseek-ai/DeepSeek-R1",#endpoint_url="https://api-inference.huggingface.co/models/cortexso/deepseek-r1:7b",
+            huggingfacehub_api_token=os.getenv("HF_TOKEN"),
+            task="text-generation",
+            #max_tokens=1024
+        )
+        self.reasoning_llm  = ChatHuggingFace(llm=llm_endpoint)
+
+        self.llm = ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                model="mistralai/Mistral-7B-Instruct-v0.3",
+                huggingfacehub_api_token=os.getenv("HF_TOKEN"),
+                task="text-generation",                
+            )
+
+        )
+
+
+
+    def ask_llm(self, question: str, do_reasoning=False)->str:
+        prompt = question
+        messages = [HumanMessage(content=prompt)]
+        response = self.llm.invoke(messages) 
+        answer = str(response.content)      
+        return answer
+
     @staticmethod
     def default_qa_function(question: str) -> str:
         """Placeholder QA function (override with your CodeAgent)"""
@@ -35,22 +74,138 @@ class GAIAAnsweringWorkflow:
         """Default GAIA formatting"""
         return answer #f"\\boxed{{{answer}}}"
     
+    def check_context_independent(self, state: AgentState)->bool:
+        if ctx := state.get("context"):
+            if ctx.get("filename"): 
+                return False
+        prompt = f"""
+        
+        I have a CodeAgent based on the text-to-text model that can use Internet search and parse the information found. 
+        If this approach is enough to successfully cope with the task, then we will call such a task an "easy question"
+
+        AS AN ERUDITE PERSON  YOU must analyze how difficult it will be to solve the next question 
+        <<{state["question"]}>>
+        
+        If you think that the question is easy, then return an empty string. Important! You should NOT add any symbols to the output in this case!
+        If the question concerns the use of additional resources such as complex analysis of downloaded files or resources on the Internet, then return an action plan
+
+        """
+        reply = self.ask_llm(prompt, True)
+        prompt = f""" The reasonings from other LLM is provided: <<{reply}>>
+        You have to Summarize: 
+        output either empty string ('') for easy question 
+        or extract action plan for non-easy question
+        """
+        reply = self.ask_llm(prompt, False)
+        if reply:
+            state["reasoning"].append(reply)
+            return False
+        return True         
+
+
     def build_workflow(self) -> Any:
         """Construct and compile the LangGraph workflow"""
         # Create graph
         workflow = StateGraph(AgentState)
         
         # Add nodes
+        workflow.add_node("preparations", self.preparations_node)
+        workflow.add_node("triage", self.triage_node)
+        workflow.add_node("deep_processing", self.deep_processing_node)
         workflow.add_node("generate_answer", self.generate_answer_node)
         workflow.add_node("format_output", self.format_output_node)
         
         # Define edges
-        workflow.set_entry_point("generate_answer")
+        workflow.set_entry_point("preparations")
+        workflow.add_edge("preparations", "triage")
+        workflow.add_conditional_edges("triage"
+                                        , self.check_context_independent
+                                        , {
+                                            True: "generate_answer",
+                                            False: "deep_processing"
+                                        })
+        workflow.add_edge("deep_processing", "format_output")
+
         workflow.add_edge("generate_answer", "format_output")
         workflow.add_edge("format_output", END)
         
         return workflow.compile()
+
+    def extract_noted_urls_with_llm(self, question: str) -> List[str]:
+        """Use LLM to extract URLs specifically noted in the question"""
+        prompt = f"""
+        Analyze the following question and extract ONLY URLs that are explicitly noted or referenced.
+        Return each URL on a separate line. If no URLs are noted, return an empty string.
+        
+        QUESTION: {question}
+        
+        Respond ONLY with the URLs, one per line, with no additional text or formatting.
+        """
+        
+        try:
+            # Use your LLM to generate the response
+            response = self.ask_llm(prompt)
+            
+            # Parse the response to extract URLs
+            urls = []
+            for line in response.split('\n'):
+                line = line.strip()
+                #if line.startswith(('http://', 'https://', 'www.')):
+                urls.append(line)
+            
+            return urls
+        except Exception as e:
+            print(f"LLM-based URL extraction failed: {str(e)}")
+            return []
+
+    def download_file(self, task_id: str, file_name: str) -> str:
+        """Download file from API and return local path"""
+        try:
+            #os.makedirs("files", exist_ok=True)
+            file_path = f"./{file_name}"#files/{file_id}"
+            api_base_url: str = "https://agents-course-unit4-scoring.hf.space"
+            api_endpoint = f"{api_base_url}/files/{task_id}"
+            response = requests.get(api_endpoint)
+            response.raise_for_status()
+            
+            with open(file_path, "wb") as f:
+                f.write(response.content)
+                
+            print(f"File saved: {file_path}")
+            return file_path
+        except Exception as e:
+            print(f"File download failed: {str(e)}")
+            return ""
+
+    def preparations_node(self, state: AgentState) -> dict:        
+        if not state["context"]:
+            return {}
+        """Node to prepare resources"""
+        context = state["context"]
+        question = state["question"]
+        uris = state["uris"]
+        
+        # 1. Handle file_id in context
+        if file_name:= context.get("file_name"):
+            file_path = self.download_file(context["task_id"], file_name)
+            if file_path:
+                uris.append(file_path)
+        
+        # 2. Extract URLs from question
+        found_urls = self.extract_noted_urls_with_llm(question)
+        if found_urls:
+            uris.extend(found_urls)
+            print(f"Added {len(found_urls)} URL(s) from question")
+        
+        return {"uris": uris}
+        
+
+    def triage_node(self, state: AgentState) -> dict:
+        return {}
     
+    def deep_processing_node(self, state: AgentState) -> dict:
+        return {}    
+
     def generate_answer_node(self, state: AgentState) -> dict:
         """Node that executes the question answering tool"""
         try:
@@ -62,13 +217,45 @@ class GAIAAnsweringWorkflow:
     
     def format_output_node(self, state: AgentState) -> dict:
         """Node that formats the answer for GAIA benchmark"""
+        prompt = f"""
+        
+        As a very smart person, you should formulate what should be the output format of the answer to the question:
+        <<{state["question"]}>>
+
+        You must formulate it very briefly and clearly!
+        The common requirements is:
+        << 
+    OUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+
+    If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise, and don't include additional text. 
+    If the answer is a number, represent it with digits.
+    If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
+    If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+    >>
+        But you have to figure out how the answer should looks like in the given case and reformulate requirement according to the specified question
+
+        """
+        format_requirement = self.ask_llm(prompt, True)
+
+        prompt = f"""
+Your attentiveness and responsibility are very much needed! We are solving a strict test that is automatically checked, so we must formulate the answer in strict accordance with the task and the required format! Even one incorrect symbol in the answer can fail the task! Pull yourself together!
+
+You will be required to produce an output of the answer, but formatted in accordance with the task
+
+Received answer: <<{state['answer']}>>
+
+Format requirements: <<{format_requirement}>>   
+
+Do NOT include << >> in your answer! Don't use full answer formulations! If you are asked about a number it MUST be just a number, nothing more! Each time it should be a clear answer (checked automatically) 
+        """ 
         try:
-            formatted = self.formatter(state["answer"])
-            return {"formatted_answer": formatted}
+            formatted = self.ask_llm(prompt)
+            return {"formatted_answer": formatted.strip()}
         except Exception as e:
             return {"formatted_answer": f"\\boxed{{\\text{{Formatting error: {str(e)}}}}}"}
     
-    def __call__(self, question: str) -> str:
+    
+    def __call__(self, question: str, context: Dict|None=None) -> str:
         """
         Execute the agent workflow for a given question
         
@@ -80,9 +267,13 @@ class GAIAAnsweringWorkflow:
         """
         # Initialize state
         initial_state = {
+            "context": context,
             "question": question,
             "answer": "",
-            "formatted_answer": ""
+            "formatted_answer": "",
+            "format_requirement": "",    
+            "uris": [],
+            "reasoning": []            
         }
         
         # Execute workflow
@@ -102,7 +293,7 @@ if __name__ == "__main__":
     # Create agent instance
     agent = GAIAAnsweringWorkflow(
         qa_function=gen_question_answer,
-        formatter=lambda ans: f"ANSWER: \\boxed{{{ans}}}"  # Custom formatting
+        formatter=lambda ans: ans #f"ANSWER: \\boxed{{{ans}}}"  # Custom formatting
     )
     
     # Test cases