Update app.py

app.py

@@ -0,0 +1,375 @@
+"""LangGraph Agent with Gradio Interface"""
+import os
+import gradio as gr
+import requests
+import pandas as pd
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.vectorstores import Chroma
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+# Load environment variables
+load_dotenv()
+
+# Tool Definitions
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers."""
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers."""
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers."""
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers."""
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers."""
+    return a % b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results."""
+    try:
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content}\n</Document>' 
+             for doc in search_docs])
+        return {"wiki_results": formatted_search_docs}
+    except Exception as e:
+        return {"wiki_results": f"Error: {str(e)}"}
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results."""
+    try:
+        search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+        formatted_search_docs = "\n\n---\n\n".join(
+            [f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content}\n</Document>' 
+             for doc in search_docs])
+        return {"web_results": formatted_search_docs}
+    except Exception as e:
+        return {"web_results": f"Error: {str(e)}"}
+
+@tool
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 results."""
+    try:
+        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content[:1000]}\n</Document>' 
+             for doc in search_docs])
+        return {"arvix_results": formatted_search_docs}
+    except Exception as e:
+        return {"arvix_results": f"Error: {str(e)}"}
+
+# System Prompt Setup
+try:
+    with open("system_prompt.txt", "r", encoding="utf-8") as f:
+        system_prompt = f.read()
+    sys_msg = SystemMessage(content=system_prompt)
+except FileNotFoundError:
+    sys_msg = SystemMessage(content="Default system prompt")
+
+# Vector Store Setup with error handling
+try:
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+    vector_store = Chroma(
+        collection_name="documents",
+        embedding_function=embeddings,
+        persist_directory="./chroma_db"
+    )
+except Exception as e:
+    print(f"Error initializing vector store: {e}")
+    vector_store = None
+
+# Tool Configuration with null check
+tools = [
+    multiply, add, subtract, divide, modulus,
+    wiki_search, web_search, arvix_search
+]
+
+if vector_store:
+    tools.append(
+        create_retriever_tool(
+            vector_store.as_retriever(),
+            name="Question Search",
+            description="Retrieves similar questions from vector store"
+        )
+    )
+else:
+    print("Warning: Vector store not initialized. Question Search tool disabled.")
+
+# Model Configuration
+MODEL_REGISTRY = {
+    "gemini-2.0-flash": {
+        "model": "gemini-2.0-flash",
+        "temperature": 0,
+        "max_tokens": 2048
+    }
+}
+
+def get_llm(model_name: str = "gemini-2.0-flash"):
+    """Initialize LLM with error handling"""
+    config = MODEL_REGISTRY.get(model_name, MODEL_REGISTRY["gemini-2.0-flash"])
+    try:
+        return ChatGoogleGenerativeAI(
+            model=config["model"],
+            temperature=config["temperature"],
+            max_tokens=config["max_tokens"]
+        )
+    except Exception as e:
+        print(f"Error initializing {model_name}: {e}")
+        return None
+
+# Updated Graph Builder Function
+def build_graph():
+    """Build LangGraph agent workflow with Gemini model"""
+    primary_llm = get_llm("gemini-2.0-flash")
+    
+    llms = [llm for llm in [primary_llm] if llm is not None]
+    
+    if not llms:
+        raise RuntimeError("Failed to initialize any LLM")
+
+    current_llm_index = 0
+
+    def assistant(state: MessagesState):
+        nonlocal current_llm_index
+        for attempt in range(len(llms)):
+            try:
+                llm = llms[current_llm_index]
+                llm_with_tools = llm.bind_tools(tools)
+                response = llm_with_tools.invoke(state["messages"])
+                current_llm_index = (current_llm_index + 1) % len(llms)  # Rotate LLMs
+                return {"messages": [response]}
+            except Exception as e:
+                print(f"Model {llms[current_llm_index].model} failed: {e}")
+                current_llm_index = (current_llm_index + 1) % len(llms)
+                if attempt == len(llms) - 1:
+                    error_msg = HumanMessage(content=f"All models failed: {str(e)}")
+                    return {"messages": [error_msg]}
+
+    def retriever(state: MessagesState):
+        try:
+            if vector_store:
+                similar_questions = vector_store.similarity_search(
+                    state["messages"][0].content,
+                    k=1
+                )
+                example_content = "Similar question reference: \n\n" + \
+                    (similar_questions[0].page_content if similar_questions 
+                     else "No similar questions found")
+            else:
+                example_content = "Vector store not available"
+            
+            return {"messages": [sys_msg] + state["messages"] + [HumanMessage(content=example_content)]}
+        except Exception as e:
+            error_msg = HumanMessage(content=f"Retrieval error: {str(e)}")
+            return {"messages": [error_msg]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+    
+    return builder.compile()
+
+
+class BasicAgent:
+    """LangGraph Agent Interface"""
+    def __init__(self):
+        self.graph = build_graph()
+
+    def __call__(self, question: str) -> str:
+        try:
+            messages = [HumanMessage(content=question)]
+            result = self.graph.invoke({"messages": messages})
+            last_message = result['messages'][-1].content
+            
+            # Improved content extraction
+            if "FINAL ANSWER: " in last_message:
+                answer_part = last_message.split("FINAL ANSWER: ")[-1].strip()
+                if answer_part.endswith('"}'):
+                    return answer_part[:-2].strip()
+                return answer_part
+            elif "Answer:" in last_message:
+                answer_part = last_message.split("Answer:")[-1].strip()
+                if answer_part.endswith('"}'):
+                    return answer_part[:-2].strip()
+                return answer_part
+            return last_message
+        except Exception as e:
+            return f"Agent processing error: {str(e)}"
+
+
+
+
+
+
+# Updated Agent Class
+# class BasicAgent:
+#     """LangGraph Agent Interface"""
+#     def __init__(self):
+#         self.graph = build_graph()
+
+#     def __call__(self, question: str) -> str:
+#         try:
+#             messages = [HumanMessage(content=question)]
+#             result = self.graph.invoke({"messages": messages})
+#             last_message = result['messages'][-1].content
+            
+#             # Improved content extraction
+#             if "FINAL ANSWER: " in last_message:
+#                 return last_message.split("FINAL ANSWER: ")[-1].strip()
+#             elif "Answer:" in last_message:
+#                 return last_message.split("Answer:")[-1].strip()
+#             return last_message
+#         except Exception as e:
+#             return f"Agent processing error: {str(e)}"
+
+
+
+
+# Gradio Interface Functions
+def run_and_submit_all(profile: gr.OAuthProfile | None):
+    """Evaluation runner function"""
+    if not profile:
+        return "Please Login to Hugging Face with the button.", None
+
+    space_id = os.getenv("SPACE_ID")
+    api_url = "https://agents-course-unit4-scoring.hf.space"
+    username = profile.username
+    results_log = []
+
+    try:
+        agent = BasicAgent()
+        agent_code = f"https://huggingface.co/spaces/ {space_id}/tree/main"
+        
+        # Fetch questions
+        response = requests.get(f"{api_url}/questions", timeout=15)
+        response.raise_for_status()
+        questions_data = response.json()
+
+        # Process questions
+        answers_payload = []
+        for item in questions_data:
+            task_id = item.get("task_id")
+            question_text = item.get("question")
+            if not task_id or not question_text:
+                continue
+                
+            try:
+                answer = agent(question_text)
+                answers_payload.append({
+                    "task_id": task_id,
+                    "submitted_answer": answer
+                })
+                results_log.append({
+                    "Task ID": task_id,
+                    "Question": question_text,
+                    "Submitted Answer": answer
+                })
+            except Exception as e:
+                results_log.append({
+                    "Task ID": task_id,
+                    "Question": question_text,
+                    "Submitted Answer": f"AGENT ERROR: {e}"
+                })
+
+        # Submit answers
+        submission_data = {
+            "username": username.strip(),
+            "agent_code": agent_code,
+            "answers": answers_payload
+        }
+        
+        response = requests.post(f"{api_url}/submit", json=submission_data, timeout=60)
+        response.raise_for_status()
+        result_data = response.json()
+        
+        final_status = (
+            f"Submission Successful!\nOverall Score: {result_data.get('score', 'N/A')}%\n"
+            f"Correct: {result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')}\n"
+            f"Message: {result_data.get('message', 'No message')}"
+        )
+        return final_status, pd.DataFrame(results_log)
+    
+    except Exception as e:
+        return f"Error: {str(e)}", pd.DataFrame(results_log)
+
+# Gradio UI Setup
+with gr.Blocks() as demo:
+    gr.Markdown("# Basic Agent Evaluation Runner")
+    gr.Markdown(
+        """
+        **Instructions:**
+        1. Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
+        2. Log in to your Hugging Face account using the button below. This uses your HF username for submission.
+        3. Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
+        ---
+        **Disclaimers:**
+        Once clicking on the "submit button, it can take quite some time (this is the time for the agent to go through all the questions).
+        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance, for the delay process of the submit button, a solution could be to cache the answers and submit in a separate action or even to answer the questions in async.
+        """
+    )
+
+    gr.LoginButton()
+
+    run_button = gr.Button("Run Evaluation & Submit All Answers")
+
+    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
+    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
+
+    run_button.click(
+        fn=run_and_submit_all,
+        outputs=[status_output, results_table]
+    )
+
+if __name__ == "__main__":
+    print("\n" + "-"*30 + " App Starting " + "-"*30)
+    # Check for SPACE_HOST and SPACE_ID at startup for information
+    space_host_startup = os.getenv("SPACE_HOST")
+    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup
+
+    if space_host_startup:
+        print(f"✅ SPACE_HOST found: {space_host_startup}")
+        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
+    else:
+        print("ℹ️ SPACE_HOST environment variable not found (running locally?).")
+
+    if space_id_startup: # Print repo URLs if SPACE_ID is found
+        print(f"✅ SPACE_ID found: {space_id_startup}")
+        print(f"   Repo URL: https://huggingface.co/spaces/ {space_id_startup}")
+        print(f"   Repo Tree URL: https://huggingface.co/spaces/ {space_id_startup}/tree/main")
+    else:
+        print("ℹ️ SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")
+
+    print("-"*(60 + len(" App Starting ")) + "\n")
+
+    print("Launching Gradio Interface for Basic Agent Evaluation...")
+    demo.launch(debug=True, share=False)