Final Submission

.gitignore

@@ -0,0 +1,3 @@
+key.txt
+*.key
+.env

app.py

@@ -9,13 +9,15 @@ from docling.document_converter import DocumentConverter
 import tqdm as notebook_tqdm
 from pydantic import BaseModel, Field
 import os
-from typing import Optional, Any, Literal, Dict, List, Tuple
+from typing import Optional, Any, Literal, Dict, List, Tuple, Type, Annotated
+from operator import add
 from typing_extensions import TypedDict
 from langgraph.graph import StateGraph, START, END
 from langgraph.types import Command
 from langchain_openai import ChatOpenAI
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import PydanticOutputParser  # Add this import
 # from langfuse.callback import CallbackHandler
 import gradio as gr
 import contextlib
@@ -26,36 +28,149 @@ import re
 from typing import Union
 from dotenv import load_dotenv
 
+# Load environment variables from .env file
 load_dotenv()
 
-# %% [markdown]
-# ## Telemetry/ Observability
-# - Used for debugging, disabled on prod
+# Use environment variables for API keys
+USE_GOOGLE = False
+API_KEY = os.environ.get("NEBIUS_KEY")
+MODEL_NAME = None
+ENDPOINT_URL = None
+
+# Try these models one by one to see which ones actually exist
+NEBIUS_MODELS = [
+    "meta-llama/Llama-2-7b-chat-hf",  # Try this first
+    "mistralai/Mistral-7B-Instruct-v0.2",  # Then this
+    "microsoft/DialoGPT-medium",  # Then this
+    "openai/gpt-3.5-turbo",  # Or this
+    "Qwen2.5-Coder-7B",  # Keep the original as fallback
+    "QwQ-32B"
+]
+
+def list_nebius_models():
+    """List all available models from Nebius API."""
+    try:
+        import requests
+        
+        headers = {
+            "Authorization": f"Bearer {API_KEY}",
+            "Content-Type": "application/json"
+        }
+        
+        # Try the models endpoint
+        response = requests.get(
+            f"{ENDPOINT_URL}models",
+            headers=headers,
+            timeout=10
+        )
+        
+        if response.status_code == 200:
+            models = response.json()
+            print("Available models:")
+            for model in models.get('data', []):
+                print(f"  - {model.get('id', 'Unknown')}")
+            return [model.get('id') for model in models.get('data', [])]
+        else:
+            print(f"Failed to fetch models: {response.status_code}")
+            print(f"Response: {response.text}")
+            return []
+            
+    except Exception as e:
+        print(f"Error fetching models: {str(e)}")
+        return []
 
-# %%
-langfuse_handler = None
-# langfuse_handler = CallbackHandler()
-if langfuse_handler:
-    TRACING = True
+def test_available_models():
+    """Test which models are actually available."""
+    
+    # First try to get the actual model list
+    available_models = list_nebius_models()
+    
+    if available_models:
+        print(f"Found {len(available_models)} models from API")
+        test_models = available_models[:6]  # Test first 6 models
+    else:
+        # Fallback to common model names that might work
+        test_models = [
+            "gpt-3.5-turbo",
+            "gpt-4",
+            "claude-3-haiku",
+            "llama-2-7b-chat",
+            "mistral-7b-instruct",
+            "qwen-7b-chat"
+        ]
+    
+    for model in test_models:
+        try:
+            print(f"Testing model: {model}")
+            global MODEL_NAME
+            MODEL_NAME = model
+            
+            # Simple test call
+            llm = ChatOpenAI(
+                model=model,
+                api_key=API_KEY,
+                base_url=ENDPOINT_URL,
+                max_completion_tokens=50,
+                timeout=10,
+                temperature=0
+            )
+            
+            response = llm.invoke("Hello")
+            print(f"✅ {model} works!")
+            return model  # Return the first working model
+            
+        except Exception as e:
+            print(f"❌ {model} failed: {str(e)}")
+            continue
+    
+    print("⚠️ No working models found")
+    return None
 
-# %% [markdown]
-# ## API Key
+# Call this function when setting up API key
+def setup_api_key(nebius_key=None, model_name=None):
+    global API_KEY, MODEL_NAME, ENDPOINT_URL, USE_GOOGLE
+    
+    # First try user-provided key (from UI)
+    if nebius_key:
+        API_KEY = nebius_key
+        ENDPOINT_URL = "https://api.studio.nebius.com/v1/"
+        
+        # Test which model actually works
+        if model_name:
+            MODEL_NAME = model_name
+        else:
+            working_model = test_available_models()
+            if working_model:
+                MODEL_NAME = working_model
+            else:
+                print("No working models found")
+                return False
+                
+        print(f"Using user-provided Nebius API key with model: {MODEL_NAME}")
+        return True
+    
+    # Next try environment variable
+    if API_KEY:
+        ENDPOINT_URL = "https://api.studio.nebius.com/v1/"
+        
+        if model_name:
+            MODEL_NAME = model_name
+        else:
+            working_model = test_available_models()
+            if working_model:
+                MODEL_NAME = working_model
+            else:
+                print("No working models found")
+                return False
+                
+        print(f"Using Nebius API from environment variable with model: {MODEL_NAME}")
+        return True
+    
+    print("No API key found. Please provide a Nebius API key.")
+    return False
 
-# %%
-USE_GOOGLE = False
-try:
-    API_KEY = os.environ["NEBIUS_KEY"]
-    MODEL_NAME = "Qwen/Qwen3-30B-A3B-fast"
-    ENDPOINT_URL = "https://api.studio.nebius.com/v1/"
-    print("Using Nebius API for Research Assistant")
-except:
-    try:
-        API_KEY = os.environ["GOOGLE_API_KEY"]
-        MODEL_NAME = os.environ["GOOGLE_DEPLOYMENT_NAME"]
-        USE_GOOGLE = True
-        print("Using Google API for Research Assistant")
-    except:
-        raise ValueError("No NEBIUS or GOOGLE API Key was found")
+# Initialize with environment variables if available
+setup_api_key()
 
 # %% [markdown]
 # ## Structured outputs
@@ -75,6 +190,20 @@ class MultiStepPlan(BaseModel):
     reasoning : str = Field("", description="The multi-step reasoning required to break down the user query in a plan.")
     plan : List[Literal["summary_agent", "synthesis_agent", "future_scope_agent", "critique_agent"]] = Field("END", description="The list of agents required to fulfill the user request determined by the Orchestrator.")
 
+class PaperSummary(BaseModel):
+    key_findings: List[str] = Field(
+        default_factory=lambda: ["No key findings available due to processing error"],
+        description="List of key findings from the paper"
+    )
+    methodology: str = Field(
+        default="Methodology not available due to processing error",
+        description="Summary of the methodology used in the paper"
+    )
+    conclusion: str = Field(
+        default="Conclusion not available due to processing error", 
+        description="Summary of the paper's conclusion"
+    )
+
 # %% [markdown]
 # ## Agent state
 
@@ -86,20 +215,19 @@ class AgentDescription(TypedDict):
     system_prompt : str
 
 class ResearchAgentState(BaseModel):
-    """State of the Academic Research Assistant."""
-    user_query : Optional[str] = Field("", description="User task for the agents to fulfill.")
-    iterations : Optional[int] = Field(0, description="Counter for iterative refinement.")
-    max_iterations : Optional[int] = Field(2, description="Maximum number of refinement iterations.")
-    available_agents : Dict[str, AgentDescription] = Field(description="A dictionary of the available agents.")
-    research_papers: List[Tuple[str, str]] = Field([], description="List of tuples containing (filename, content) of uploaded research papers.")
-    phase : Literal["PLAN", "EXECUTE", "ANSWER"] = Field("PLAN", description="Current phase of the agent")
-    messages : List[Tuple[str,str]] = Field([], description="List of agent thoughts (agent, agent response).")
-    final_answer : str = Field("", description="Final answer generated after task execution.")
-    plan : List[Literal["summary_agent", "synthesis_agent", "future_scope_agent", "critique_agent"]] = Field([], description="The current list of tasks to execute")
-    synthesis_of_findings: Optional[str] = Field("", description="A synthesized summary of findings from all provided papers.")
-    identified_gaps: Optional[str] = Field("", description="A consolidated list of research gaps.")
-    future_directions_report: Optional[str] = Field("", description="The final report on future research directions.")
-    critique: str = Field("", description="Written feedback from the critique agent regarding the generated report.")
+    """State for the research agent."""
+    research_papers: Annotated[List[Tuple[str, str]], add] = Field(default_factory=list)  # List of (filename, content) tuples
+    summary: Annotated[List[Dict], add] = Field(default_factory=list)  # List of paper summaries
+    user_query: str = Field(default="")  # Remove annotation - only set once
+    phase: str = Field(default="PLAN")  # PLAN, EXECUTE, ANSWER
+    plan: List[str] = Field(default_factory=list)  # List of agent names to call in order
+    messages: Annotated[List[Tuple[str, str]], add] = Field(default_factory=list)  # List of (agent, message) tuples
+    critique: Optional[str] = Field(default=None)  # Optional critique of the analysis
+    available_agents: Dict[str, Dict] = Field(default_factory=dict)  # Mapping of agent name to agent description
+    final_answer: Optional[str] = Field(default=None)  # Final answer to the user's query
+    max_iterations: int = Field(default=1)  # Maximum number of iterations for processing
+    synthesis_of_findings: Optional[str] = Field(default=None)  # Remove Annotated - only set once
+    future_directions_report: Optional[str] = Field(default=None)  # Remove Annotated - only set once
 
 # %% [markdown]
 # ## System prompts
@@ -220,48 +348,94 @@ def read_file_content(file: Union[str, Path]) -> str:
     else:
         return ""
 
-def call_llm(system_prompt, user_prompt, response_format : Any = None) -> Any:
-    """
-    Call LLM with provided system prompt and user prompt and the response format that should be enforced
-    """
+# %% [markdown]
+# ## LLM call
 
-    if USE_GOOGLE:
-        llm = ChatGoogleGenerativeAI(
-            model = MODEL_NAME,
-            google_api_key = API_KEY,
-            temperature = 0,
-            max_tokens = None,
-            timeout = None,
-            max_retries = 2
-        )
-    else:
-        llm = ChatOpenAI(
-        model=MODEL_NAME,
-        api_key=API_KEY,
-        base_url=ENDPOINT_URL,
-        max_completion_tokens=None,
-        timeout=60,
-        max_retries=0,
-        temperature=0
-        )
+# %%
+def call_llm(system_prompt, user_prompt, response_format=None):
+    """Call LLM with system and user prompt, optionally parsing to a specific format"""
+    global API_KEY, MODEL_NAME, ENDPOINT_URL
     
-    if response_format is not None:  
-        llm = llm.with_structured_output(response_format)
-
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", "{system_prompt}"),
-        ("user", "{user_prompt}")
-    ])
+    if not API_KEY:
+        print("Error: API key is not set")
+        # Return a default instance for the response_format class
+        if response_format and hasattr(response_format, "__name__"):
+            try:
+                if response_format.__name__ == "MultiStepPlan":
+                    return MultiStepPlan(
+                        reasoning="Error occurred: API key not set",
+                        plan=["summary_agent", "synthesis_agent", "future_scope_agent"]
+                    )
+                elif response_format.__name__ == "PaperSummary":
+                    return PaperSummary()  # Uses default values from Field definitions
+                else:
+                    # Generic attempt to create an instance with default values
+                    return response_format()
+            except Exception as e:
+                print(f"Failed to create default instance: {str(e)}")
+        return None
+        
+    try:
+        if USE_GOOGLE:
+            llm = ChatGoogleGenerativeAI(
+                model=MODEL_NAME,
+                google_api_key=API_KEY,
+                temperature=0
+            )
+        else:
+            llm = ChatOpenAI(
+                model=MODEL_NAME,
+                api_key=API_KEY,
+                base_url=ENDPOINT_URL,
+                max_completion_tokens=None,
+                timeout=60,
+                max_retries=2,
+                temperature=0
+            )
+            
+        if response_format is not None:  
+            llm = llm.with_structured_output(response_format)
 
-    chain = prompt | llm
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", "{system_prompt}"),
+            ("user", "{user_prompt}")
+        ])
 
-    response = chain.invoke({
-        "system_prompt":system_prompt,
-        "user_prompt": user_prompt
-    })
+        chain = prompt | llm
 
-    return response 
+        print(f"Calling model: {MODEL_NAME}")
+        response = chain.invoke({
+            "system_prompt": system_prompt,
+            "user_prompt": user_prompt
+        })
 
+        return response
+        
+    except Exception as e:
+        print(f"Error in call_llm: {str(e)}")
+        if hasattr(e, 'response') and hasattr(e.response, 'json'):
+            try:
+                error_details = e.response.json()
+                print(f"API Error details: {error_details}")
+            except:
+                pass
+                
+        # Create a default response based on the response_format class
+        if response_format and hasattr(response_format, "__name__"):
+            try:
+                if response_format.__name__ == "MultiStepPlan":
+                    return MultiStepPlan(
+                        reasoning="Error occurred while calling the LLM API. Using default plan.",
+                        plan=["summary_agent", "synthesis_agent", "future_scope_agent"]
+                    )
+                elif response_format.__name__ == "PaperSummary":
+                    return PaperSummary()  # Uses default values from Field definitions
+                else:
+                    # Generic attempt to create an instance with default values
+                    return response_format()
+            except Exception as e:
+                print(f"Failed to create default instance: {str(e)}")
+        return None
 def serialize_messages(messages : List[Tuple[str,str]]) -> str:
     "Returns a formatted message history of previous messages"
     return "\n" +"\n".join(f"**{role}:**\n{content}" for role, content in messages)
@@ -302,9 +476,11 @@ def orchestrator_agent(state: ResearchAgentState) -> Command:
     """Central orchestration logic to determine the next agent to call."""
     
     if not state.research_papers:
-        state.final_answer = "### ❗️ The research assistant needs at least one research paper to begin.\n" \
-        "👈🏽 Please upload one or more research papers in the '📚 Research Materials' tab."  
-        return Command(goto=END, update=state)
+        return Command(
+            goto=END, 
+            update={"final_answer": "### ❗️ The research assistant needs at least one research paper to begin.\n" \
+                   "👈🏽 Please upload one or more research papers in the '📚 Research Materials' tab."}
+        )
     
     if state.phase == "PLAN":
         agent_descriptions = "\n".join([
@@ -317,93 +493,216 @@ def orchestrator_agent(state: ResearchAgentState) -> Command:
         **Agents:**
         {agent_descriptions}
         
-        Based on the user's query, create a logical sequence of agents to call. For example, to find future scope, you should first summarize the papers, then synthesize them, and finally call the future_scope_agent.
+        Based on the user's query, create a logical sequence of agents to call. For example, to find future scope, you should first summarize the papers, then synthesize them, and then call the future_scope_agent.
         
         **IMPORTANT:** Always include the summary_agent as the first step when working with research papers. Every task requires proper paper summaries before analysis can begin.
-        
-        **[EXAMPLE]**
-        **USER QUERY:** Summarize these papers and tell me what to research next.
-        **REASONING:** The user wants a summary and future research directions. I need to first run the summary_agent on all papers, then the synthesis_agent, and finally the future_scope_agent.
-        **PLAN:** ["summary_agent", "synthesis_agent", "future_scope_agent"]
         """
 
         user_prompt = state.user_query
-        state.messages.append(("user_query", state.user_query))
 
         response = call_llm(system_prompt, user_prompt, MultiStepPlan)
         
-        # Enforce summary_agent as the first step if not already included
-        if not response.plan or response.plan[0] != "summary_agent":
-            print("⚠️ Enforcing summary_agent as first step in the plan")
-            response.plan.insert(0, "summary_agent")
+        # Handle None response by providing a default plan
+        if response is None:
+            print("⚠️ Failed to get response from LLM. Using default plan.")
+            plan = ["summary_agent", "synthesis_agent", "future_scope_agent"]
+            print("="*40)
+            print("🤖 DEFAULT ORCHESTRATOR PLAN (LLM call failed)")
+            print("="*40)
+            print("\n📝 Reasoning: Default plan due to LLM call failure\n")
+            print("🔗 Planned Steps:")
+            for i, step in enumerate(plan, 1):
+                print(f"  {i}. {step}")
+            print("="*40)
+            print("⚙️ EXECUTE PLAN")
+            print("="*40 + "\n")
+            
+            # Create update dict that only modifies necessary fields
+            updates = {
+                "plan": plan,
+                "phase": "EXECUTE"
+            }
+            
+            # Only add user_query to messages if it's not already there
+            if not any(msg[0] == "user_query" for msg in state.messages):
+                updates["messages"] = [("user_query", state.user_query)]
+            
+            return Command(goto=plan[0], update=updates)
         
-        print("="*40)
-        print("🤖 ORCHESTRATOR PLAN")
-        print("="*40)
-        print(f"\n📝 Reasoning:\n{response.reasoning}\n")
-        print("🔗 Planned Steps:")
-        for i, step in enumerate(response.plan, 1):
-            print(f"  {i}. {step}")
-        print("="*40)
-        print("⚙️ EXECUTE PLAN")
-        print("="*40 + "\n")
-        state.plan = response.plan
-        state.phase = "EXECUTE"
+        # If response exists but plan is empty, use default plan
+        try:
+            # Enforce summary_agent as the first step if not already included
+            if not hasattr(response, 'plan') or not response.plan:
+                print("⚠️ Response from LLM did not contain a valid plan. Using default plan.")
+                response.plan = ["summary_agent", "synthesis_agent", "future_scope_agent"]
+            elif response.plan[0] != "summary_agent":
+                print("⚠️ Enforcing summary_agent as first step in the plan")
+                response.plan.insert(0, "summary_agent")
+            
+            print("="*40)
+            print("🤖 ORCHESTRATOR PLAN")
+            print("="*40)
+            print(f"\n📝 Reasoning:\n{getattr(response, 'reasoning', 'No reasoning provided')}\n")
+            print("🔗 Planned Steps:")
+            for i, step in enumerate(response.plan, 1):
+                print(f"  {i}. {step}")
+            print("="*40)
+            print("⚙️ EXECUTE PLAN")
+            print("="*40 + "\n")
+            
+            # Create update dict that only modifies necessary fields
+            updates = {
+                "plan": response.plan,
+                "phase": "EXECUTE"
+            }
+            
+            # Only add user_query to messages if it's not already there
+            if not any(msg[0] == "user_query" for msg in state.messages):
+                updates["messages"] = [("user_query", state.user_query)]
+            
+            return Command(goto=response.plan[0], update=updates)
+            
+        except Exception as e:
+            # Final fallback if response processing fails
+            print(f"⚠️ Error processing LLM response: {str(e)}. Using default plan.")
+            plan = ["summary_agent", "synthesis_agent", "future_scope_agent"]
+            
+            # Create update dict that only modifies necessary fields
+            updates = {
+                "plan": plan,
+                "phase": "EXECUTE"
+            }
+            
+            # Only add user_query to messages if it's not already there
+            if not any(msg[0] == "user_query" for msg in state.messages):
+                updates["messages"] = [("user_query", state.user_query)]
+            
+            return Command(goto=plan[0], update=updates)
     
     if len(state.plan) == 0 and state.phase == "EXECUTE":
-        state.phase = "ANSWER"
         return Command(
             goto="final_answer_tool",
-            update=state
+            update={"phase": "ANSWER"}
         )
+        
     if state.phase == "EXECUTE":
-        agent = state.plan.pop(0)
+        next_agent = state.plan[0]
+        remaining_plan = state.plan[1:]
         return Command(
-            goto=agent,
-            update=state)
+            goto=next_agent,
+            update={"plan": remaining_plan}
+        )
 
     if state.phase == "ANSWER":
-        state.phase = "PLAN"
-        state.messages = [("orchestrator_agent", f"\n{state.final_answer}")]
+        return Command(
+            goto=END, 
+            update={
+                "phase": "PLAN",
+                "messages": [("orchestrator_agent", f"\n{state.final_answer}")]
+            }
+        )
         
-    return Command(goto=END, update=state)
+    return Command(goto=END, update={})
 
 # %% [markdown]
 # ### Research Agents
 
 # %%
 def summary_agent(state : ResearchAgentState) -> Command:
-    "Summarizes each research paper provided."
-    
-    agent_description = state.available_agents.get("summary_agent", {})
-    system_prompt = agent_description.get("system_prompt")
+    """Creates concise, structured summaries of research papers."""
     
-    summaries = []
-    print(f"The summary agent is processing {len(state.research_papers)} paper(s) 📚")
-    for filename, content in state.research_papers:
-        user_prompt = f"[RESEARCH PAPER: {filename}]\n{content}\n[END RESEARCH PAPER]"
+    if not state.summary:
+        # Initialize empty summaries
+        print("The summary agent is processing the papers... 📝")
+        research_findings = []
         
-        print(f"  - Summarizing {filename}...")
-        response = call_llm(system_prompt, user_prompt, ResearchSummary)
-        summary_text = f"### Summary for {filename}\n\n**Key Findings:**\n"
-        for finding in response.key_findings:
-            summary_text += f"- {finding}\n"
-        summary_text += f"\n**Methodology:**\n{response.methodology}\n"
-        summary_text += f"\n**Limitations:**\n"
-        for limitation in response.limitations:
-            summary_text += f"- {limitation}\n"
-
-        summaries.append(summary_text)
-
-    agent_contribution = ("summary_agent", "\n\n---\n\n".join(summaries))
-    state.messages.append(agent_contribution)
-    print("The summary agent has finished.")
-    
-    return Command(goto="orchestrator_agent", update=state)
-
-
+        for filename, content in state.research_papers:
+            # Create a prompt for each paper
+            system_prompt = f"""You are a research summarization expert. Please read the provided research paper content and create a clear, concise, and structured summary.
+            Focus on extracting key findings, methodology, and conclusions.
+            """
+            user_prompt = f"""
+            Paper: {filename}
+            
+            Content:
+            {content[:5000]}  # Use first 5000 chars to avoid context limits
+            
+            Please provide a structured summary with key findings, methodology, and conclusions.
+            """
+            
+            response = call_llm(system_prompt, user_prompt, PaperSummary)
+            
+            # Check if we got a valid response
+            if response is None:
+                print(f"⚠️ Failed to summarize paper {filename}. Creating default summary.")
+                # Create a default summary
+                finding = {
+                    "title": filename,
+                    "key_findings": ["Error: Could not summarize this paper due to API issues."],
+                    "methodology": "Not available due to API error",
+                    "conclusion": "Not available due to API error",
+                    "source": filename
+                }
+                research_findings.append(finding)
+            else:
+                try:
+                    # Extract the key findings from the response
+                    finding = {
+                        "title": filename,
+                        "key_findings": response.key_findings if hasattr(response, 'key_findings') else ["No key findings extracted"],
+                        "methodology": response.methodology if hasattr(response, 'methodology') else "Not provided",
+                        "conclusion": response.conclusion if hasattr(response, 'conclusion') else "Not provided",
+                        "source": filename
+                    }
+                    research_findings.append(finding)
+                except Exception as e:
+                    print(f"⚠️ Error processing summary for {filename}: {str(e)}")
+                    finding = {
+                        "title": filename,
+                        "key_findings": ["Error processing paper summary."],
+                        "methodology": "Error in processing",
+                        "conclusion": "Error in processing",
+                        "source": filename
+                    }
+                    research_findings.append(finding)
+        
+        print("Paper summaries complete.")
+        
+        # Add the summaries to the message history
+        formatted_summaries = []
+        for paper in research_findings:
+            findings_text = "\n".join([f"- {finding}" for finding in paper['key_findings']])
+            formatted_summary = f"""
+        ## {paper['title']}
+        
+        ### Key Findings:
+        {findings_text}
+        
+        ### Methodology:
+        {paper['methodology']}
+        
+        ### Conclusion:
+        {paper['conclusion']}
+        """
+            formatted_summaries.append(formatted_summary)
+        
+        combined_summary = "\n\n".join(formatted_summaries)
+        
+        agent_contribution = ("summary_agent", combined_summary)
+        
+        # Return updates for both summary and messages
+        return Command(
+            goto="orchestrator_agent", 
+            update={
+                "summary": research_findings,
+                "messages": [agent_contribution]
+            }
+        )
+    else:
+        # Summaries already exist, just proceed
+        return Command(goto="orchestrator_agent", update=state)
 def synthesis_agent(state : ResearchAgentState) -> Command:
-    "Synthesizes the summaries into a cohesive narrative."
+    """Synthesizes the summaries into a cohesive narrative."""
     
     agent_description = state.available_agents.get("synthesis_agent", {})
     system_prompt = agent_description.get("system_prompt")
@@ -412,17 +711,27 @@ def synthesis_agent(state : ResearchAgentState) -> Command:
     user_prompt = f"Please synthesize the following research summaries:\n{previous_messages}"
     
     print("The synthesis agent is creating a literature review...")
-    response = call_llm(system_prompt, user_prompt).content
+    response = call_llm(system_prompt, user_prompt)
+    
+    # Handle None response
+    if response is None:
+        response_text = "Error: Could not generate synthesis due to API issues."
+        print("⚠️ Synthesis agent failed - using default response")
+    else:
+        response_text = response.content if hasattr(response, 'content') else str(response)
+    
     print("Synthesis complete.")
 
-    state.synthesis_of_findings = response
-    agent_contribution = ("synthesis_agent", response)
-    state.messages.append(agent_contribution)
-    
-    return Command(goto="orchestrator_agent", update=state)
+    # Only update messages, don't update synthesis_of_findings
+    return Command(
+        goto="orchestrator_agent", 
+        update={
+            "messages": [("synthesis_agent", response_text)]
+        }
+    )
 
 def future_scope_agent(state : ResearchAgentState) -> Command:
-    "Identifies research gaps and suggests future work."
+    """Identifies research gaps and suggests future work."""
     
     agent_description = state.available_agents.get("future_scope_agent", {})
     system_prompt = agent_description.get("system_prompt")
@@ -433,25 +742,35 @@ def future_scope_agent(state : ResearchAgentState) -> Command:
     print("The future scope agent is identifying research gaps...")
     response = call_llm(system_prompt, user_prompt, FutureScope)
     
-    report_text = "### Identified Research Gaps\n"
-    for gap in response.identified_gaps:
-        report_text += f"- {gap}\n"
-    report_text += "\n### Suggested Future Directions\n"
-    for direction in response.suggested_directions:
-        report_text += f"- {direction}\n"
-    report_text += f"\n### Concluding Synthesis\n{response.synthesis}"
+    # Handle None response
+    if response is None:
+        print("⚠️ Future scope agent failed - using default response")
+        report_text = "### Identified Research Gaps\n- Error: Could not identify gaps due to API issues.\n\n### Suggested Future Directions\n- Error: Could not suggest directions due to API issues.\n\n### Concluding Synthesis\nError: Could not generate synthesis due to API issues."
+    else:
+        try:
+            report_text = "### Identified Research Gaps\n"
+            for gap in response.identified_gaps:
+                report_text += f"- {gap}\n"
+            report_text += "\n### Suggested Future Directions\n"
+            for direction in response.suggested_directions:
+                report_text += f"- {direction}\n"
+            report_text += f"\n### Concluding Synthesis\n{response.synthesis}"
+        except Exception as e:
+            print(f"⚠️ Error processing future scope response: {str(e)}")
+            report_text = "### Error\nCould not process future scope analysis due to response format issues."
 
     print("Future scope analysis complete.")
     
-    state.future_directions_report = report_text
-    agent_contribution = ("future_scope_agent", report_text)
-    state.messages.append(agent_contribution)
-
-    return Command(goto="orchestrator_agent", update=state)
-
+    # Only update messages, don't update future_directions_report
+    return Command(
+        goto="orchestrator_agent", 
+        update={
+            "messages": [("future_scope_agent", report_text)]
+        }
+    )
 
 def critique_agent(state: ResearchAgentState) -> Command:
-    "Provides feedback on the generated analysis."
+    """Provides feedback on the generated analysis."""
 
     agent_description = state.available_agents.get("critique_agent", {})
     system_prompt = agent_description.get("system_prompt")
@@ -460,16 +779,25 @@ def critique_agent(state: ResearchAgentState) -> Command:
     user_prompt = f"Please critique the following research analysis:\n{previous_messages}"
     
     print("The critique agent is reviewing the analysis... 🔎")
-    response = call_llm(system_prompt, user_prompt).content
+    response = call_llm(system_prompt, user_prompt)
+    
+    # Handle None response
+    if response is None:
+        response_text = "Error: Could not generate critique due to API issues."
+        print("⚠️ Critique agent failed - using default response")
+    else:
+        response_text = response.content if hasattr(response, 'content') else str(response)
+    
     print("Critique complete.")
     
-    state.critique = response
-    agent_contribution = ("critique_agent", response)
-    state.messages.append(agent_contribution)
-
-    return Command(goto="orchestrator_agent", update=state)
-
-
+    # Only update the fields that need updating - avoid updating user_query
+    return Command(
+        goto="orchestrator_agent", 
+        update={
+            "critique": response_text,
+            "messages": [("critique_agent", response_text)]
+        }
+    )
 def final_answer_tool(state : ResearchAgentState) -> Command[Literal["orchestrator_agent"]]:
     "Final answer tool is invoked to formulate a final answer based on the agent message history"
 
@@ -492,34 +820,62 @@ def final_answer_tool(state : ResearchAgentState) -> Command[Literal["orchestrat
     Compile the final, comprehensive answer for the user based on the history.
     """
 
-    final_answer = call_llm(system_prompt, user_prompt).content
+    response = call_llm(system_prompt, user_prompt)
+    
+    # Handle None response
+    if response is None:
+        final_answer = "Error: Could not generate final answer due to API issues. Please check the logs and try again."
+        print("⚠️ Final answer tool failed - using default response")
+    else:
+        final_answer = response.content if hasattr(response, 'content') else str(response)
 
     if isinstance(final_answer, str):
         final_answer = strip_think_blocks(final_answer)
-        state.final_answer = final_answer
     
+    # Only update the final_answer field, not the entire state
     return Command(
         goto="orchestrator_agent",
-        update=state
+        update={"final_answer": final_answer}
     )
 
 
 # %% [markdown]
-# ## Build Graph
+# ## Graph Definition
 
 # %%
-builder = StateGraph(ResearchAgentState)
-
-builder.add_node(orchestrator_agent)
-builder.add_node(summary_agent)
-builder.add_node(synthesis_agent)
-builder.add_node(future_scope_agent)
-builder.add_node(critique_agent)
-builder.add_node(final_answer_tool)
+def init_state():
+    """Initialize the state with default values."""
+    return ResearchAgentState(available_agents=available_agents)
+
+graph = StateGraph(ResearchAgentState)
+graph.add_node("orchestrator_agent", orchestrator_agent)
+graph.add_node("summary_agent", summary_agent)
+graph.add_node("synthesis_agent", synthesis_agent)
+graph.add_node("future_scope_agent", future_scope_agent)
+graph.add_node("critique_agent", critique_agent)
+graph.add_node("final_answer_tool", final_answer_tool)
+
+# Define the edges
+graph.add_edge(START, "orchestrator_agent")
+
+# Fix the parameter name from 'router' to the correct parameter name
+graph.add_conditional_edges(
+    "orchestrator_agent",
+    lambda state: (
+        state.plan[0] if state.phase == "EXECUTE" and state.plan 
+        else "final_answer_tool" if state.phase == "ANSWER"
+        else END
+    )
+)
 
-builder.add_edge(START, "orchestrator_agent")
+graph.add_edge("summary_agent", "orchestrator_agent")
+graph.add_edge("synthesis_agent", "orchestrator_agent")
+graph.add_edge("future_scope_agent", "orchestrator_agent")
+graph.add_edge("critique_agent", "orchestrator_agent")
+graph.add_edge("final_answer_tool", "orchestrator_agent")
 
-graph = builder.compile()
+# Compile the graph
+graph = graph.compile()
 
 # %% [markdown]
 # ## Gradio functions
@@ -530,31 +886,70 @@ def extract_research_papers(
     paper_files,
     max_iterations: int
 ) -> tuple[str, Dict, bool]:
-    """
-    Run the extraction pipeline and return output logs + state as a dict.
-    """
-    output_text = ""
-    if not paper_files:
-        return "Please upload at least one paper.", state_dict, False
-
-    papers = []
-    try:
-        for file in paper_files:
-            # Gradio uploads files to a temp directory, file.name gives the full path
-            content = read_file_content(file.name)
-            papers.append((os.path.basename(file.name), content))
-        output_text += f"Successfully processed {len(papers)} paper(s)."
-    except Exception as e:
-        return f"Reading input files failed: {str(e)}", state_dict, False
-
-    state = ResearchAgentState.model_validate(state_dict)
-    state.research_papers = papers
+    """Extract text from research papers and update state."""
+    
+    # Create a new ResearchAgentState or update existing one
+    if isinstance(state_dict, dict):
+        state = ResearchAgentState(**state_dict)
+    else:
+        state = ResearchAgentState()
+    
+    # Set max_iterations safely
     state.max_iterations = max_iterations
     
-    state_dict = type_conversion(state, ResearchAgentState)
+    if not paper_files:
+        return "Please upload at least one research paper to analyze.", state.model_dump(), False
     
-    return output_text, state_dict, True
-
+    console_output = StringIO()
+    with contextlib.redirect_stdout(console_output):
+        papers = []
+        
+        for file in paper_files:
+            try:
+                filename = file.name.split("/")[-1]
+                print(f"📄 Processing {filename}...")
+                
+                if filename.lower().endswith(".pdf"):
+                    # Fix DocumentConverter usage - it likely uses a different method name
+                    try:
+                        converter = DocumentConverter()
+                        # Try different method names that might exist
+                        if hasattr(converter, 'pdf_to_text'):
+                            content = converter.pdf_to_text(file.name)
+                        elif hasattr(converter, 'extract_text'):
+                            content = converter.extract_text(file.name)
+                        else:
+                            # Fallback to PyPDF2 if available
+                            import PyPDF2
+                            content = ""
+                            with open(file.name, 'rb') as pdf_file:
+                                pdf_reader = PyPDF2.PdfReader(pdf_file)
+                                for page_num in range(len(pdf_reader.pages)):
+                                    content += pdf_reader.pages[page_num].extract_text()
+                    except ImportError:
+                        print("⚠️ PDF conversion libraries not available. Please install PyPDF2.")
+                        continue
+                        
+                elif filename.lower().endswith(".docx"):
+                    doc = docx.Document(file.name)
+                    content = "\n".join([p.text for p in doc.paragraphs])
+                elif filename.lower().endswith((".txt", ".md")):
+                    with open(file.name, "r") as f:
+                        content = f.read()
+                else:
+                    print(f"⚠️ Unsupported file format: {filename}")
+                    continue
+                
+                papers.append((filename, content))
+                print(f"✅ Successfully extracted {len(content)} characters from {filename}")
+                
+            except Exception as e:
+                print(f"❌ Error processing {file.name}: {str(e)}")
+        
+        state.research_papers = papers
+        print(f"📊 Extracted content from {len(papers)} files.")
+        
+    return console_output.getvalue(), state.model_dump(), len(papers) > 0
 def call_orchestrator(state_dict : Dict, user_query : str):
     "Function prototype to call the orchestrator agent"
     state = ResearchAgentState.model_validate(state_dict)
@@ -562,10 +957,17 @@ def call_orchestrator(state_dict : Dict, user_query : str):
     state.user_query = user_query
     buffer = StringIO()
     with contextlib.redirect_stdout(buffer):
-        if langfuse_handler:  # Use langfuse_handler instead of TRACING
-            result = graph.invoke(input=state, config={"callbacks": [langfuse_handler]})
-        else:
-           result = graph.invoke(input=state)
+        config = {}  # Use empty config for now
+        
+        try:
+            result = graph.invoke(input=state, config=config)
+            output_text = buffer.getvalue()
+            result_dict = type_conversion(result, ResearchAgentState)
+            return output_text, result_dict, True
+        except Exception as e:
+            error_msg = f"An error occurred during processing: {str(e)}"
+            output_text = buffer.getvalue() + "\n" + error_msg
+            return output_text, state_dict, False
 
     result_dict = type_conversion(result, ResearchAgentState)
 
@@ -593,8 +995,80 @@ with gr.Blocks() as research_assistant_server:
 
     state_dict = gr.State(value=ResearchAgentState(available_agents=available_agents).model_dump())
     extraction_successful = gr.State(value=False)
+    api_key_set = gr.State(value=API_KEY is not None)
     
     with gr.Tabs():
+        with gr.TabItem("🔑 API Key Setup"):
+            gr.Markdown("### Set up your Nebius API Key")
+            gr.Markdown("A valid API key is required to use this research assistant. You can either provide it here or set it as an environment variable.")
+            
+            with gr.Row():
+                nebius_key_input = gr.Textbox(
+                    label="Nebius API Key", 
+                    placeholder="Enter your Nebius API key here...",
+                    type="password",
+                    value=""
+                )
+            
+            # Add model discovery section
+            with gr.Row():
+                discover_models_button = gr.Button("🔍 Discover Available Models", variant="secondary")
+                test_model_input = gr.Textbox(
+                    label="Or manually test a model name:",
+                    placeholder="e.g., gpt-3.5-turbo"
+                )
+            
+            available_models_display = gr.Textbox(
+                label="Available Models",
+                lines=5,
+                interactive=False
+            )
+            
+            with gr.Row():
+                model_dropdown = gr.Dropdown(
+                    choices=NEBIUS_MODELS,
+                    value=MODEL_NAME or NEBIUS_MODELS[0],
+                    label="Select Nebius Model",
+                    allow_custom_value=True
+                )
+            
+            api_key_status = gr.Markdown("⚠️ **No API key detected.** Please enter your Nebius API key." if API_KEY is None else "✅ **API key configured.** You're ready to use the assistant.")
+            
+            save_key_button = gr.Button("Save API Key", variant="primary")
+            
+            def discover_models(key):
+                if not key:
+                    return "Please enter an API key first."
+                
+                global API_KEY, ENDPOINT_URL
+                API_KEY = key
+                ENDPOINT_URL = "https://api.studio.nebius.com/v1/"
+                
+                models = list_nebius_models()
+                if models:
+                    return "Available models:\n" + "\n".join([f"- {model}" for model in models])
+                else:
+                    return "Could not fetch models. Please check your API key."
+            
+            discover_models_button.click(
+                fn=discover_models,
+                inputs=[nebius_key_input],
+                outputs=[available_models_display]
+            )
+            
+            def save_api_key(key, model):
+                success = setup_api_key(key, model)
+                if success:
+                    return f"✅ **API key saved successfully!** Using model: {MODEL_NAME}", True
+                else:
+                    return "❌ **Invalid API key.** Please check and try again.", False
+            
+            save_key_button.click(
+                fn=save_api_key,
+                inputs=[nebius_key_input, model_dropdown],
+                outputs=[api_key_status, api_key_set]
+            )
+            
         with gr.TabItem("📚 Research Materials"):
             gr.Markdown("### 🍉 Feed the assistant with the research papers you want to analyze.")
 
@@ -612,13 +1086,20 @@ with gr.Blocks() as research_assistant_server:
             extract_button = gr.Button("Process Papers", variant="primary")
 
             extract_console_output = gr.Textbox(label="Logs / Console Output")
+            
+            # Modify extract_research_papers to check for API key
+            def extract_with_api_check(state_dict, paper_files, max_iterations, api_key_set):
+                if not api_key_set:
+                    return "⚠️ Please set up your API key in the 'API Key Setup' tab first.", state_dict, False
+                return extract_research_papers(state_dict, paper_files, max_iterations)
 
             extract_button.click(
-                fn=extract_research_papers,
-                inputs=[state_dict, research_papers_files, max_iterations],
+                fn=extract_with_api_check,
+                inputs=[state_dict, research_papers_files, max_iterations, api_key_set],
                 outputs=[extract_console_output, state_dict, extraction_successful]
             )
 
+        # Rest of your tabs remain the same, but with API key checks for Q&A
         with gr.TabItem("🤖 Q&A Chatbot"):
             examples = """ℹ️ **Example Queries**
 - Summarize the key findings from these papers.
@@ -628,28 +1109,51 @@ with gr.Blocks() as research_assistant_server:
             gr.Markdown(examples)
             user_query = gr.Textbox(label="Ask your research question", value="Identify the main gaps and suggest future work.", interactive=True)
             button = gr.Button("Ask the Research Assistant 🔬🧠", variant="primary")
-            qa_orchestrator_completed = gr.State(value=False)
-
-            @gr.render(inputs=[qa_orchestrator_completed,state_dict])
-            def show_qa_results(qa_flag, state_dict):
-                if qa_flag: 
-                    gr.Markdown(state_dict.get("final_answer", "❗️ No final answer provided: please check the logs."))
-                else:
-                    gr.Markdown("### Upload papers and ask a question to get started.")
+            
+            # Replace the @gr.render with a proper output textbox
+            qa_output = gr.Markdown(
+                label="Research Assistant Response",
+                value="### 📝 Upload papers and ask a question to get started.",
+                elem_id="qa_output"
+            )
             
             output_logs = gr.Textbox(label="Logs/ Console Output", lines=10)
 
-            def reset_elements(qa_flag : bool, output_logs : str) -> bool:
-                return False, "Generating response..."
+            def call_with_api_check(state_dict, user_query, api_key_set):
+                """Wrapper to check API key before calling orchestrator."""
+                if not API_KEY:
+                    error_msg = "⚠️ Please set up your API key in the 'API Key Setup' tab first."
+                    return error_msg, error_msg, state_dict
+                
+                if not state_dict.get("research_papers"):
+                    error_msg = "### ❗️ No Research Papers Found\n\n👈🏽 Please upload research papers in the '📚 Research Materials' tab first."
+                    return error_msg, error_msg, state_dict
+                
+                try:
+                    logs, updated_state, success = call_orchestrator(state_dict, user_query)
+                    
+                    if success and updated_state.get("final_answer"):
+                        final_answer = updated_state.get("final_answer")
+                        return final_answer, logs, updated_state
+                    else:
+                        error_msg = f"### ❗️ Processing Failed\n\n{logs}\n\nPlease check the logs above for details."
+                        return error_msg, logs, state_dict
+                        
+                except Exception as e:
+                    error_msg = f"### ❗️ An Error Occurred\n\n```\n{str(e)}\n```\n\nPlease check your API key and try again."
+                    return error_msg, f"Error: {str(e)}", state_dict
+
+            def reset_output():
+                """Reset the output when starting a new query."""
+                return "### 🤖 Processing your request...\n\nPlease wait while the research assistant analyzes your papers and generates a response.", "Generating response..."
 
             button.click(
-                fn=reset_elements,
-                inputs=[qa_orchestrator_completed, output_logs],
-                outputs=[qa_orchestrator_completed, output_logs]
+                fn=reset_output,
+                outputs=[qa_output, output_logs]
             ).then(
-                fn=call_orchestrator,
-                inputs=[state_dict, user_query],
-                outputs=[output_logs, state_dict, qa_orchestrator_completed]
+                fn=call_with_api_check,
+                inputs=[state_dict, user_query, api_key_set],
+                outputs=[qa_output, output_logs, state_dict]
             )
 
         with gr.TabItem("🔎 What's under the hood?"):

key.txt

@@ -1,2 +0,0 @@
-eyJhbGciOiJIUzI1NiIsImtpZCI6IlV6SXJWd1h0dnprLVRvdzlLZWstc0M1akptWXBvX1VaVkxUZlpnMDRlOFUiLCJ0eXAiOiJKV1QifQ.eyJzdWIiOiJnaXRodWJ8Mzc1MzgyNTIiLCJzY29wZSI6Im9wZW5pZCBvZmZsaW5lX2FjY2VzcyIsImlzcyI6ImFwaV9rZXlfaXNzdWVyIiwiYXVkIjpbImh0dHBzOi8vbmViaXVzLWluZmVyZW5jZS5ldS5hdXRoMC5jb20vYXBpL3YyLyJdLCJleHAiOjE5MDcyMjQ2NTcsInV1aWQiOiI0NGQyOGU3ZC0xMjRmLTQ1ZjgtYTczMS0yNWRmY2Q1NTkyZTgiLCJuYW1lIjoiTkVCSVVTX0tFWSIsImV4cGlyZXNfYXQiOiIyMDMwLTA2LTA5VDA4OjM3OjM3KzAwMDAifQ.djouDFHucL8mm3NVfDU0VTfsPjuDGx7LcOzWRu8isl8
-eyJhbGciOiJIUzI1NiIsImtpZCI6IlV6SXJWd1h0dnprLVRvdzlLZWstc0M1akptWXBvX1VaVkxUZlpnMDRlOFUiLCJ0eXAiOiJKV1QifQ.eyJzdWIiOiJnaXRodWJ8Mzc1MzgyNTIiLCJzY29wZSI6Im9wZW5pZCBvZmZsaW5lX2FjY2VzcyIsImlzcyI6ImFwaV9rZXlfaXNzdWVyIiwiYXVkIjpbImh0dHBzOi8vbmViaXVzLWluZmVyZW5jZS5ldS5hdXRoMC5jb20vYXBpL3YyLyJdLCJleHAiOjE5MDcyNDI2NzYsInV1aWQiOiIxY2NkMDBlMy01NmNkLTRlMGMtODE3MS1lZmQ5MWQ2NmI5YWMiLCJuYW1lIjoiTkVCSVVTX0tFWSIsImV4cGlyZXNfYXQiOiIyMDMwLTA2LTA5VDEzOjM3OjU2KzAwMDAifQ.Olq0zvLp8X-F9fmOKdcZjX0WpeHwmURQOagTQJt5Mts