lang_graph_agent.py

import os
import requests
from typing import TypedDict, List, Dict, Any, Optional, Annotated
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.utils.function_calling import convert_to_openai_function
from duckduckgo_search import DDGS

class AgentState(TypedDict):
    task_id: str
    question: str
    file_contents: Optional[str]
    file_link: Optional[str]
    messages: List[Any]
    search_results: Optional[List[Dict[str, str]]]
    final_answer: Optional[str]

def read_file(task_id: str) -> str:
    """Read a file from the Agent Evaluation API using the task ID"""
    api_url = "https://agents-course-unit4-scoring.hf.space"
    response = requests.get(f"{api_url}/file/{task_id}")
    if response.status_code == 200:
        return response.text
    else:
        raise Exception(f"Failed to read file: {response.status_code} - {response.text}")

def web_search(query: str, num_results: int = 3) -> List[Dict[str, str]]:
    """Perform a web search using DuckDuckGo"""
    with DDGS() as ddgs:
        results = []
        try:
            for r in ddgs.text(query, max_results=num_results):
                # Handle potential variations in response structure
                title = r.get('title', '')
                link = r.get('url', r.get('link', ''))  # Try both 'url' and 'link' keys
                body = r.get('body', r.get('snippet', ''))  # Try both 'body' and 'snippet' keys
                
                results.append({
                    'title': title,
                    'link': link,
                    'snippet': body
                })
        except Exception as e:
            print(f"Error during web search: {e}")
            # Return empty results in case of error
            results = [{
                'title': 'Search Error',
                'link': '',
                'snippet': f'Failed to perform search: {str(e)}'
            }]
        return results

class LangGraphAgent:
    def __init__(self):
        print("LangGraphAgent initialized.")
        self.llm = ChatOpenAI(model="o1")
        # self.llm = ChatOpenAI(model="gpt-4.1", temperature=0)
        self.graph = self._build_graph()

    def _build_graph(self) -> StateGraph:
        workflow = StateGraph(AgentState)

        # Define the read task node
        workflow.add_node("read_task", self.read_task_step)

        # Define the file link node
        workflow.add_node("create_file_link", self.create_file_link)

        # Define the search node
        workflow.add_node("search", self.search_step)
        
        # Define nodes
        workflow.add_node("generate_answer", self.generate_answer)
        workflow.add_node("verify_answer", self.verify_answer)
        workflow.add_node("verify_format", self.verify_format)

        # Connect the nodes
        # workflow.set_entry_point("read_task")
        # workflow.add_edge("read_task", "create_file_link")
        workflow.set_entry_point("create_file_link")
        workflow.add_edge("create_file_link", "search")
        workflow.add_edge("search", "generate_answer")
        workflow.add_edge("generate_answer", "verify_answer")
        workflow.add_edge("verify_answer", "verify_format")
        workflow.set_finish_point("verify_format")

        return workflow

    def read_task_step(self, state: AgentState) -> AgentState:
        """Read the task file and store both the contents and question"""
        task_content = read_file(state['task_id'])
        state['file_contents'] = task_content
        state['question'] = task_content
        return state

    def create_file_link(self, state: AgentState) -> AgentState:
        """Create a link to the file using the task_id"""
        api_url = "https://agents-course-unit4-scoring.hf.space"
        state['file_link'] = f"{api_url}/file/{state['task_id']}"
        return state

    def search_step(self, state: AgentState) -> AgentState:
        """Perform web search based on the question"""
        search_results = web_search(state['question'])
        state['search_results'] = search_results
        return state

    def generate_answer(self, state: AgentState) -> AgentState:
        """Generate final answer using search results"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful AI assistant that provides accurate answers based on search results and file contents."),
            ("human", "Question:\n{question}\n\nFile Link:\n{file_link}\n\nSearch Results:\n{search_results}\n\nPlease provide just the answer based on these inputs. Do not include any additional text."),
        ])

        # Format search results for the prompt
        formatted_results = '\n'.join([f"Title: {r['title']}\nSnippet: {r['snippet']}\nLink: {r['link']}\n" 
                                     for r in state['search_results']])

        # Generate response
        response = self.llm.invoke(
            prompt.format_messages(
                question=state['question'],
                file_link=state['file_link'],
                search_results=formatted_results
            )
        )

        state['final_answer'] = response.content
        return state

    def verify_answer(self, state: AgentState) -> AgentState:
        """Verify that the generated answer is correct based on search results"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a verification assistant. Your job is to verify if the provided answer is correct based on both the original task file and search results."),
            ("human", "Question:\n{question}\n\nFile Link:\n{file_link}\n\nSearch Results:\n{search_results}\n\nProposed Answer: {answer}\n\nVerify if this answer is correct based on both the task file and search results. If it's incorrect or unsupported, provide a corrected answer. If it's correct, return the same answer."),
        ])

        formatted_results = '\n'.join([f"Title: {r['title']}\nSnippet: {r['snippet']}\nLink: {r['link']}\n" 
                                     for r in state['search_results']])

        response = self.llm.invoke(
            prompt.format_messages(
                question=state['question'],
                file_link=state['file_link'],
                search_results=formatted_results,
                answer=state['final_answer']
            )
        )

        state['final_answer'] = response.content
        return state

    def verify_format(self, state: AgentState) -> AgentState:
        """Verify that the answer contains just the necessary words without additional text"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a formatting assistant. Your job is to ensure the answer contains only the exact information for the answer without any additional text. You will be provided with the original answer and you must return just the answer without any additional text, explanations, or formatting. Answers that are a number should be returned as just a number without any additional text."),
            ("human", "Original answer: {answer}\n\nPlease provide just the essential answer without any additional text, explanations, or formatting. Answers that are a number should be returned as just a number without any additional text."),
        ])

        response = self.llm.invoke(
            prompt.format_messages(answer=state['final_answer'])
        )

        state['final_answer'] = response.content
        return state

    def __call__(self, question: str, task_id: str) -> str:
        print(f"Agent received question: {question}")
        print(f"Task ID: {task_id}")
        
        # Initialize state
        state = {
            'question': question,
            'task_id': task_id,
            'file_link': None,
            'messages': [],
            'search_results': None,
            'final_answer': None
        }

        # Run the graph
        app = self.graph.compile()
        final_state = app.invoke(state)
        
        return final_state['final_answer']