Update main.py

main.py

@@ -1,139 +1,96 @@
 """
 GenAI Advanced Agent - Production Ready
-Architecture: FastAPI + LangGraph + Ollama + DuckDuckGo
-Features: Streaming, Memory, Tools, Structured Output, Error Handling
+Fixed: Infinite loop prevention, better tool error handling
 """
 
 import os
 import logging
 import asyncio
-import nest_asyncio  # ADD THIS
+import nest_asyncio
 from typing import Annotated, TypedDict, List, Dict, Any, Optional, AsyncGenerator
 from contextlib import asynccontextmanager
 from datetime import datetime
 from enum import Enum
 
-from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field, validator
 
-# Async & Network
 import httpx
 from duckduckgo_search import DDGS
 from bs4 import BeautifulSoup
 
-# LangChain / AI Core
 from langchain_ollama import ChatOllama
-from langchain_core.messages import HumanMessage, SystemMessage, BaseMessage, ToolMessage
+from langchain_core.messages import HumanMessage, SystemMessage, BaseMessage, ToolMessage, AIMessage
 from langchain_core.tools import tool, BaseTool
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableConfig
 
-# LangGraph
 from langgraph.graph import StateGraph, END, START
 from langgraph.prebuilt import ToolNode, tools_condition
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.checkpoint.base import BaseCheckpointSaver
 
-# FIX: Apply nest_asyncio for Jupyter/IPython environments
 nest_asyncio.apply()
 
 # --------------------------------------------------------------------------------------
-# CONFIGURATION & LOGGING
+# CONFIGURATION
 # --------------------------------------------------------------------------------------
 
 class Settings(BaseModel):
-    """Application configuration"""
     MODEL_NAME: str = "qwen2.5:3b"
     BASE_URL: str = "http://localhost:11434"
     TEMPERATURE: float = 0.3
-    MAX_TOKENS: int = 4096
-    TIMEOUT: float = 30.0
+    MAX_ITERATIONS: int = 3  # Prevent infinite loops
     MAX_SEARCH_RESULTS: int = 5
     MAX_CONTENT_LENGTH: int = 4000
-    LOG_LEVEL: str = "INFO"
+    TIMEOUT: float = 30.0
     
     class Config:
         env_file = ".env"
 
 settings = Settings()
 
-# Structured Logging
 logging.basicConfig(
-    level=getattr(logging, settings.LOG_LEVEL),
-    format='%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger("GenAI-Agent")
 
 # --------------------------------------------------------------------------------------
-# MODELS & SCHEMAS
+# MODELS
 # --------------------------------------------------------------------------------------
 
-class MessageType(str, Enum):
-    HUMAN = "human"
-    AI = "ai"
-    SYSTEM = "system"
-    TOOL = "tool"
-
-class ChatMessage(BaseModel):
-    role: MessageType
-    content: str
-    timestamp: Optional[datetime] = Field(default_factory=datetime.now)
-    metadata: Optional[Dict[str, Any]] = None
-
 class ChatRequest(BaseModel):
-    query: str = Field(..., min_length=1, max_length=10000, description="User query")
-    thread_id: str = Field(..., min_length=1, description="Conversation thread ID")
-    stream: bool = Field(default=True, description="Enable streaming response")
-    context: Optional[List[ChatMessage]] = Field(default=None, description="Previous messages")
+    query: str = Field(..., min_length=1, max_length=10000)
+    thread_id: str = Field(..., min_length=1)
+    stream: bool = Field(default=True)
     
     @validator('thread_id')
     def validate_thread_id(cls, v):
-        if not v.strip():
-            raise ValueError("thread_id cannot be empty")
         return v.strip()
 
-class ChatResponse(BaseModel):
-    response: str
-    thread_id: str
-    tools_used: List[str]
-    tokens_used: Optional[int] = None
-    processing_time: float
-
-class HealthStatus(BaseModel):
-    status: str
-    model: str
-    version: str
-    timestamp: datetime
-
 # --------------------------------------------------------------------------------------
-# STATE MANAGEMENT
+# STATE MANAGEMENT - FIXED: Added iteration counter
 # --------------------------------------------------------------------------------------
 
 class AgentState(TypedDict):
-    """LangGraph state definition"""
     messages: Annotated[List[BaseMessage], "add_messages"]
     thread_id: str
     tools_used: Annotated[List[str], "append"]
-    metadata: Dict[str, Any]
+    iteration_count: int  # NEW: Track iterations to prevent loops
+    last_tool_result: Optional[str]  # NEW: Track last tool result
 
 # --------------------------------------------------------------------------------------
-# TOOLS IMPLEMENTATION
+# TOOLS - FIXED: Better error messages and validation
 # --------------------------------------------------------------------------------------
 
 class ToolRegistry:
-    """Centralized tool management with caching and metrics"""
-    
     def __init__(self):
         self._tools: Dict[str, BaseTool] = {}
-        self._metrics: Dict[str, Dict] = {}
         
     def register(self, tool_instance: BaseTool):
         self._tools[tool_instance.name] = tool_instance
-        self._metrics[tool_instance.name] = {"calls": 0, "errors": 0, "avg_time": 0}
         return tool_instance
     
     def get(self, name: str) -> Optional[BaseTool]:
@@ -141,159 +98,116 @@ class ToolRegistry:
     
     def all_tools(self) -> List[BaseTool]:
         return list(self._tools.values())
-    
-    def record_usage(self, name: str, duration: float, error: bool = False):
-        if name in self._metrics:
-            self._metrics[name]["calls"] += 1
-            if error:
-                self._metrics[name]["errors"] += 1
-            # Update running average
-            prev_avg = self._metrics[name]["avg_time"]
-            n = self._metrics[name]["calls"]
-            self._metrics[name]["avg_time"] = (prev_avg * (n-1) + duration) / n
 
 tool_registry = ToolRegistry()
 
 @tool
 async def web_search(query: str, max_results: int = 5) -> str:
     """
-    Advanced web search with result ranking and filtering.
-    Use for: current events, technical documentation, news, facts verification.
+    Search the web for current information. Returns formatted search results.
+    Use this for: current events, documentation, news, facts.
     """
-    start_time = asyncio.get_event_loop().time()
+    if not query or len(query.strip()) < 2:
+        return "ERROR: Query too short or empty"
     
     def _sync_search(q: str):
         try:
             with DDGS() as ddgs:
-                results = ddgs.text(q, max_results=max_results)
-                return list(results)
+                results = list(ddgs.text(q, max_results=max_results))
+                return results
         except Exception as e:
-            logger.error(f"Search failed: {e}")
-            raise
+            logger.error(f"DDGS Error: {e}")
+            return f"ERROR: Search failed - {str(e)}"
     
     try:
-        logger.info(f"🔍 Web search: {query}")
+        logger.info(f"🔍 Searching: {query}")
         results = await asyncio.to_thread(_sync_search, query)
         
-        if not results:
-            return "No relevant results found."
+        # Handle error string return
+        if isinstance(results, str) and results.startswith("ERROR"):
+            return results
+            
+        if not results or len(results) == 0:
+            return "ERROR: No results found for this query. Try a different search term."
         
         formatted = []
         for idx, r in enumerate(results, 1):
-            formatted.append(
-                f"[{idx}] {r.get('title', 'Untitled')}\n"
-                f"URL: {r.get('href', 'N/A')}\n"
-                f"Summary: {r.get('body', 'No description')}\n"
-            )
-        
-        duration = asyncio.get_event_loop().time() - start_time
-        tool_registry.record_usage("web_search", duration)
+            title = r.get('title', 'Untitled')
+            link = r.get('href', 'N/A')
+            body = r.get('body', 'No description')
+            formatted.append(f"[{idx}] {title}\nURL: {link}\nSummary: {body}\n")
         
         return "\n".join(formatted)
         
     except Exception as e:
-        tool_registry.record_usage("web_search", 0, error=True)
-        return f"Search error: {str(e)}"
+        logger.error(f"Search error: {e}")
+        return f"ERROR: {str(e)}"
 
 @tool
-async def read_webpage(url: str, extract_code: bool = False) -> str:
+async def read_webpage(url: str) -> str:
     """
-    Intelligent webpage reader with content extraction and cleaning.
-    Use for: deep technical details, documentation, code examples.
+    Read content from a specific URL. Use for detailed documentation.
     """
-    start_time = asyncio.get_event_loop().time()
+    if not url.startswith(('http://', 'https://')):
+        return "ERROR: Invalid URL format"
     
     try:
         logger.info(f"📖 Reading: {url}")
         
         headers = {
-            "User-Agent": "Mozilla/5.0 (AppleWebKit/537.36) GenAI-Agent/2.0",
-            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
-            "Accept-Language": "en-US,en;q=0.5",
-            "Accept-Encoding": "gzip, deflate",
-            "DNT": "1",
-            "Connection": "keep-alive",
+            "User-Agent": "Mozilla/5.0 GenAI-Agent/2.0",
+            "Accept": "text/html,application/xhtml+xml",
         }
         
-        async with httpx.AsyncClient(timeout=20.0, follow_redirects=True) as client:
+        async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
             response = await client.get(url, headers=headers)
             response.raise_for_status()
         
         soup = BeautifulSoup(response.text, 'lxml')
         
         # Remove noise
-        for element in soup(["script", "style", "nav", "footer", "header", 
-                            "aside", "advertisement", "svg", "iframe"]):
+        for element in soup(["script", "style", "nav", "footer", "header"]):
             element.decompose()
         
-        # Extract main content (prefer article/main tags)
         main_content = soup.find('article') or soup.find('main') or soup.find('body')
-        
-        if extract_code:
-            # Extract code blocks specifically
-            code_blocks = main_content.find_all(['pre', 'code'])
-            code_content = '\n\n'.join(
-                block.get_text() for block in code_blocks if block.get_text().strip()
-            )
-            if code_content:
-                return f"Code extracted:\n{code_content[:settings.MAX_CONTENT_LENGTH]}"
+        text = main_content.get_text(separator='\n') if main_content else soup.get_text()
         
         # Clean text
-        text = main_content.get_text(separator='\n') if main_content else soup.get_text()
         lines = (line.strip() for line in text.splitlines())
         chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
         clean_text = '\n'.join(chunk for chunk in chunks if chunk)
         
-        # Smart truncation with context preservation
-        if len(clean_text) > settings.MAX_CONTENT_LENGTH:
-            truncated = clean_text[:settings.MAX_CONTENT_LENGTH]
-            # Try to end at a sentence boundary
-            last_period = truncated.rfind('.')
-            if last_period > len(truncated) * 0.8:
-                truncated = truncated[:last_period + 1]
-            clean_text = truncated + "\n\n[Content truncated...]"
-        
-        duration = asyncio.get_event_loop().time() - start_time
-        tool_registry.record_usage("read_webpage", duration)
-        
-        return clean_text
+        if len(clean_text) < 100:
+            return "ERROR: Content too short or page blocked"
+            
+        return clean_text[:settings.MAX_CONTENT_LENGTH] + "\n[Content truncated...]" if len(clean_text) > settings.MAX_CONTENT_LENGTH else clean_text
         
     except httpx.HTTPStatusError as e:
-        tool_registry.record_usage("read_webpage", 0, error=True)
-        return f"HTTP Error {e.response.status_code}: Unable to access {url}"
+        return f"ERROR: HTTP {e.response.status_code} - Unable to access page"
     except Exception as e:
-        tool_registry.record_usage("read_webpage", 0, error=True)
-        return f"Scraping error: {str(e)}"
+        return f"ERROR: {str(e)}"
 
 @tool
 async def calculate(expression: str) -> str:
     """
-    Safe mathematical expression evaluator.
-    Use for: calculations, data processing, unit conversions.
+    Calculate mathematical expressions safely.
     """
     try:
-        # Safe eval with limited scope
-        allowed_names = {
-            "abs": abs, "round": round, "max": max, "min": min,
-            "sum": sum, "pow": pow, "len": len
-        }
-        result = eval(expression, {"__builtins__": {}}, allowed_names)
+        allowed = {"abs": abs, "round": round, "max": max, "min": min, "sum": sum, "pow": pow}
+        result = eval(expression, {"__builtins__": {}}, allowed)
         return f"Result: {result}"
     except Exception as e:
-        return f"Calculation error: {str(e)}"
+        return f"ERROR: Invalid expression - {str(e)}"
 
-# Register all tools
 tool_registry.register(web_search)
 tool_registry.register(read_webpage)
 tool_registry.register(calculate)
 
 # --------------------------------------------------------------------------------------
-# LANGGRAPH AGENT ARCHITECTURE
+# LANGGRAPH - FIXED: Added iteration limit and better routing
 # --------------------------------------------------------------------------------------
 
 class AgentBuilder:
-    """Factory for building configurable LangGraph agents"""
-    
     def __init__(self, model_name: str, base_url: str, temperature: float = 0.3):
         self.model_name = model_name
         self.base_url = base_url
@@ -306,7 +220,6 @@ class AgentBuilder:
         return self
         
     def build(self) -> StateGraph:
-        # Initialize LLM with tools
         llm = ChatOllama(
             model=self.model_name,
             base_url=self.base_url,
@@ -315,59 +228,122 @@ class AgentBuilder:
             num_ctx=8192
         ).bind_tools(self.tools)
         
-        # System prompt with dynamic tool descriptions
         tool_descriptions = "\n".join([
             f"- {t.name}: {t.description}" for t in self.tools
         ])
         
-        system_prompt = f"""You are an advanced GenAI technical assistant with access to real-time tools.
+        # FIXED: Stronger instructions to prevent loops
+        system_prompt = f"""You are an advanced AI assistant with tools.
 
-AVAILABLE TOOLS:
+TOOLS AVAILABLE:
 {tool_descriptions}
 
-CORE INSTRUCTIONS:
-1. **Always use tools** for current information, calculations, or external data
-2. **Chain tools intelligently**: Search → Read → Analyze
-3. **Format responses**:
-   - Start with a brief executive summary
-   - Use markdown headers (##) for sections
-   - For code, use XML tags: <code lang="python">your code</code>
-   - Cite sources when using web data
-4. **Be concise** but thorough. Avoid hallucinations.
+CRITICAL RULES:
+1. **MAXIMUM 2 tool calls per conversation** - After that, answer with available info
+2. **NEVER call the same tool twice** with similar queries
+3. **If a tool returns ERROR**, do NOT retry - explain the limitation to user
+4. **If web_search returns no results**, tell user you couldn't find info online
+5. **DO NOT LOOP** - If you've searched once, don't search again
+
+Response Format:
+- Start with brief summary
+- Use ## for headers
+- Use XML for code: <code lang="python">code</code>
 
 Current date: {datetime.now().strftime("%Y-%m-%d")}
 """
         
-        # Agent node
         async def agent_node(state: AgentState):
+            # FIXED: Check iteration limit
+            if state.get("iteration_count", 0) >= settings.MAX_ITERATIONS:
+                logger.warning("Max iterations reached, forcing end")
+                # Force final response
+                messages = state["messages"] + [
+                    AIMessage(content="I've reached the maximum number of tool calls. Let me provide the best answer based on the information gathered so far.")
+                ]
+                return {
+                    "messages": messages,
+                    "iteration_count": state["iteration_count"],
+                    "tools_used": state.get("tools_used", []),
+                    "last_tool_result": state.get("last_tool_result")
+                }
+            
             messages = [SystemMessage(content=system_prompt)] + state["messages"]
             response = await llm.ainvoke(messages)
             
             # Track tool usage
+            tools_used = state.get("tools_used", []).copy()
             if response.tool_calls:
-                state["tools_used"].extend([
-                    tc["name"] for tc in response.tool_calls
-                ])
+                tools_used.extend([tc["name"] for tc in response.tool_calls])
             
-            return {"messages": [response], "tools_used": state["tools_used"]}
+            return {
+                "messages": [response],
+                "iteration_count": state.get("iteration_count", 0) + 1,
+                "tools_used": tools_used,
+                "last_tool_result": state.get("last_tool_result")
+            }
         
-        # Tool node with error handling
-        tool_node = ToolNode(self.tools)
+        # FIXED: Custom tool node with error tracking
+        async def tool_node_with_tracking(state: AgentState):
+            tool_node = ToolNode(self.tools)
+            result = await tool_node.ainvoke(state)
+            
+            # Check if tool returned error
+            last_msg = result["messages"][-1] if result["messages"] else None
+            if last_msg and hasattr(last_msg, 'content'):
+                content = str(last_msg.content)
+                if content.startswith("ERROR") or "No results found" in content:
+                    logger.warning(f"Tool error detected: {content[:100]}")
+                    # Add error context to state
+                    result["last_tool_result"] = "error"
+            
+            # Increment iteration count
+            result["iteration_count"] = state.get("iteration_count", 0) + 1
+            result["tools_used"] = state.get("tools_used", [])
+            
+            return result
+        
+        # FIXED: Better conditional routing
+        def should_continue(state: AgentState) -> str:
+            last_message = state["messages"][-1] if state["messages"] else None
+            
+            # Check iteration limit
+            if state.get("iteration_count", 0) >= settings.MAX_ITERATIONS:
+                logger.info("Max iterations reached, ending")
+                return END
+            
+            # Check if last tool had error
+            if state.get("last_tool_result") == "error":
+                logger.info("Previous tool had error, ending to prevent loop")
+                return END
+            
+            # Check if there are tool calls
+            if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
+                # Check if same tool being called repeatedly
+                current_tools = [tc["name"] for tc in last_message.tool_calls]
+                previous_tools = state.get("tools_used", [])
+                
+                # If web_search called more than once, prevent loop
+                if current_tools.count("web_search") > 0 and previous_tools.count("web_search") >= 1:
+                    logger.warning("Preventing web_search loop")
+                    return END
+                    
+                return "tools"
+            
+            return END
         
-        # Build graph
         workflow = StateGraph(AgentState)
         workflow.add_node("agent", agent_node)
-        workflow.add_node("tools", tool_node)
+        workflow.add_node("tools", tool_node_with_tracking)
         
         workflow.add_edge(START, "agent")
         workflow.add_conditional_edges(
             "agent",
-            tools_condition,
+            should_continue,
             {"tools": "tools", END: END}
         )
         workflow.add_edge("tools", "agent")
         
-        # Compile with memory
         if self.checkpointer:
             return workflow.compile(checkpointer=self.checkpointer)
         return workflow.compile()
@@ -376,7 +352,6 @@ Current date: {datetime.now().strftime("%Y-%m-%d")}
 # FASTAPI APPLICATION
 # --------------------------------------------------------------------------------------
 
-# Global state
 class AppState:
     def __init__(self):
         self.http_client: Optional[httpx.AsyncClient] = None
@@ -387,8 +362,6 @@ app_state = AppState()
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Application lifecycle management"""
-    # Startup
     logger.info("🚀 Starting GenAI Agent...")
     
     app_state.http_client = httpx.AsyncClient(
@@ -396,7 +369,6 @@ async def lifespan(app: FastAPI):
         limits=httpx.Limits(max_keepalive_connections=20, max_connections=100)
     )
     
-    # Initialize memory and agent
     app_state.memory = MemorySaver()
     builder = AgentBuilder(
         model_name=settings.MODEL_NAME,
@@ -405,25 +377,19 @@ async def lifespan(app: FastAPI):
     )
     app_state.agent = builder.with_memory(app_state.memory).build()
     
-    logger.info(f"✅ Agent ready with model: {settings.MODEL_NAME}")
+    logger.info(f"✅ Agent ready: {settings.MODEL_NAME}")
     yield
     
-    # Shutdown
     logger.info("🛑 Shutting down...")
     if app_state.http_client:
         await app_state.http_client.aclose()
 
-# Create FastAPI app
 app = FastAPI(
-    title="GenAI Advanced Agent API",
-    description="Production-ready AI agent with web search, RAG, and memory",
-    version="3.0.0",
-    lifespan=lifespan,
-    docs_url="/docs",
-    redoc_url="/redoc"
+    title="GenAI Agent API",
+    version="3.1.0",
+    lifespan=lifespan
 )
 
-# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -432,45 +398,28 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Security
-security = HTTPBearer(auto_error=False)
-
 # --------------------------------------------------------------------------------------
 # API ENDPOINTS
 # --------------------------------------------------------------------------------------
 
-@app.get("/health", response_model=HealthStatus)
+@app.get("/health")
 async def health_check():
-    """Health check endpoint"""
-    return HealthStatus(
-        status="healthy",
-        model=settings.MODEL_NAME,
-        version="3.0.0",
-        timestamp=datetime.now()
-    )
-
-@app.get("/tools")
-async def list_tools():
-    """List available tools and their metrics"""
     return {
-        "tools": [
-            {
-                "name": name,
-                "description": tool_registry.get(name).description,
-                "metrics": tool_registry._metrics.get(name, {})
-            }
-            for name in tool_registry._tools.keys()
-        ]
+        "status": "healthy",
+        "model": settings.MODEL_NAME,
+        "version": "3.1.0",
+        "max_iterations": settings.MAX_ITERATIONS,
+        "timestamp": datetime.now()
     }
 
 async def stream_response(query: str, thread_id: str) -> AsyncGenerator[str, None]:
-    """Generate streaming response with real-time updates"""
     config = RunnableConfig(configurable={"thread_id": thread_id})
     inputs = {
         "messages": [HumanMessage(content=query)],
         "thread_id": thread_id,
         "tools_used": [],
-        "metadata": {}
+        "iteration_count": 0,  # Initialize counter
+        "last_tool_result": None
     }
     
     yield f"event: start\ndata: {thread_id}\n\n"
@@ -489,8 +438,10 @@ async def stream_response(query: str, thread_id: str) -> AsyncGenerator[str, Non
                 yield f"event: tool_start\ndata: {tool_name}\n\n"
                 
             elif event_type == "on_tool_end":
-                output = str(event["data"].get("output", ""))[:200]
-                yield f"event: tool_end\ndata: {output}...\n\n"
+                output = str(event["data"].get("output", ""))
+                # Truncate long outputs
+                preview = output[:200] + "..." if len(output) > 200 else output
+                yield f"event: tool_end\ndata: {preview}\n\n"
                 
         yield "event: complete\ndata: done\n\n"
         
@@ -500,9 +451,6 @@ async def stream_response(query: str, thread_id: str) -> AsyncGenerator[str, Non
 
 @app.post("/chat")
 async def chat_endpoint(request: ChatRequest):
-    """
-    Main chat endpoint with streaming support
-    """
     try:
         if request.stream:
             return StreamingResponse(
@@ -511,78 +459,35 @@ async def chat_endpoint(request: ChatRequest):
                 headers={
                     "Cache-Control": "no-cache",
                     "Connection": "keep-alive",
-                    "X-Thread-ID": request.thread_id
                 }
             )
         else:
-            # Non-streaming response
             config = RunnableConfig(configurable={"thread_id": request.thread_id})
             inputs = {
                 "messages": [HumanMessage(content=request.query)],
                 "thread_id": request.thread_id,
                 "tools_used": [],
-                "metadata": {}
+                "iteration_count": 0,
+                "last_tool_result": None
             }
             
             result = await app_state.agent.ainvoke(inputs, config=config)
             final_message = result["messages"][-1]
             
-            return ChatResponse(
-                response=final_message.content,
-                thread_id=request.thread_id,
-                tools_used=result.get("tools_used", []),
-                processing_time=0.0  # Calculate if needed
-            )
+            return {
+                "response": final_message.content,
+                "thread_id": request.thread_id,
+                "tools_used": result.get("tools_used", []),
+                "iterations": result.get("iteration_count", 0)
+            }
             
     except Exception as e:
         logger.error(f"Chat error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/chat/sync", response_model=ChatResponse)
-async def chat_sync(request: ChatRequest):
-    """Synchronous chat endpoint for simple requests"""
-    return await chat_endpoint(request)
-
-@app.delete("/memory/{thread_id}")
-async def clear_memory(thread_id: str):
-    """Clear conversation memory for a thread"""
-    try:
-        # MemorySaver specific implementation
-        if hasattr(app_state.memory, 'delete'):
-            await app_state.memory.delete(thread_id)
-        return {"status": "success", "message": f"Memory cleared for {thread_id}"}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/memory/{thread_id}")
-async def get_conversation(thread_id: str):
-    """Retrieve conversation history"""
-    try:
-        config = RunnableConfig(configurable={"thread_id": thread_id})
-        # This depends on your checkpointer implementation
-        return {"thread_id": thread_id, "history": []}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# --------------------------------------------------------------------------------------
-# MAIN ENTRY - FIXED FOR JUPYTER
-# --------------------------------------------------------------------------------------
-
 def run_server():
-    """Run server with proper async handling for Jupyter"""
     import uvicorn
-    
-    # Use this instead of asyncio.run()
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=8000,
-        log_level="info"
-    )
-
-# For Jupyter/IPython - run directly
-# run_server()
+    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
 
-# For normal Python execution
 if __name__ == "__main__":
     run_server()