docrag/main.py

#!/usr/bin/env python3
"""
DocRAG - OpenAI-Compatible RAG Server

This application presents itself as a standard OpenAI API server that can be used
with any OpenAI-compatible client (like Open WebUI). Behind the scenes, it:
1. Detects URLs in user messages and auto-downloads websites
2. Ingests website content into the RAG knowledge base
3. Retrieves relevant context from the knowledge base
4. Passes the enriched context to OpenRouter for response generation

The user sees a normal chat experience, but the system is actually doing
sophisticated RAG operations in the background.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import sys
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator, Optional

# Load environment variables from .env file (look in script directory)
from dotenv import load_dotenv
SCRIPT_DIR = Path(__file__).parent.resolve()
ENV_FILE = SCRIPT_DIR / ".env"
load_dotenv(ENV_FILE)

# Also try loading from current working directory
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# FastAPI imports
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field

# Import RAG components
from rag import RAGSystem, get_rag_system
from rag.document_processor import DocumentProcessor

# Import tools
from tools import ToolManager, get_tool_manager

# Import OpenAI client for OpenRouter
from openai import AsyncOpenAI


# =============================================================================
# Configuration
# =============================================================================

class Config:
    """Application configuration from environment variables."""

    # Server settings
    HOST: str = os.getenv("HOST", "0.0.0.0")
    PORT: int = int(os.getenv("PORT", "8000"))
    DEBUG: bool = os.getenv("DEBUG", "false").lower() == "true"

    # Model settings
    MODEL_NAME: str = os.getenv("MODEL_NAME", "DocRAG")
    UPSTREAM_MODEL: str = os.getenv("UPSTREAM_MODEL", "openrouter/free")

    # OpenRouter API settings
    OPENROUTER_API_KEY: str = os.getenv("OPENROUTER_API_KEY", "")
    OPENROUTER_BASE_URL: str = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")

    # RAG settings
    EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
    VECTOR_STORE_PATH: str = os.getenv("VECTOR_STORE_PATH", "./data/vectors")
    DOCUMENTS_PATH: str = os.getenv("DOCUMENTS_PATH", "./data/documents")
    CHUNK_SIZE: int = int(os.getenv("CHUNK_SIZE", "1000"))
    CHUNK_OVERLAP: int = int(os.getenv("CHUNK_OVERLAP", "200"))
    TOP_K_RESULTS: int = int(os.getenv("TOP_K_RESULTS", "5"))

    # Tool settings
    ENABLE_TOOLS: bool = os.getenv("ENABLE_TOOLS", "true").lower() == "true"
    MAX_TOOL_ITERATIONS: int = int(os.getenv("MAX_TOOL_ITERATIONS", "5"))


config = Config()


# =============================================================================
# OpenAI-Compatible Models
# =============================================================================

class ChatMessage(BaseModel):
    """OpenAI chat message format."""
    role: str
    content: Optional[str] = None
    name: Optional[str] = None
    tool_calls: Optional[list[dict]] = None
    tool_call_id: Optional[str] = None


class ChatCompletionRequest(BaseModel):
    """OpenAI chat completion request format."""
    model: str = "DocRAG"
    messages: list[ChatMessage]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 1.0
    n: Optional[int] = 1
    stream: Optional[bool] = False
    stop: Optional[list[str]] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = 0.0
    frequency_penalty: Optional[float] = 0.0
    user: Optional[str] = None
    tools: Optional[list[dict]] = None
    tool_choice: Optional[str | dict] = None


class ChatCompletionChoice(BaseModel):
    """OpenAI chat completion choice."""
    index: int = 0
    message: ChatMessage
    finish_reason: str = "stop"


class ChatCompletionUsage(BaseModel):
    """Token usage statistics."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionResponse(BaseModel):
    """OpenAI chat completion response format."""
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = config.MODEL_NAME
    choices: list[ChatCompletionChoice]
    usage: ChatCompletionUsage = ChatCompletionUsage()


class ModelInfo(BaseModel):
    """OpenAI model info format."""
    id: str
    object: str = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "organization"


class ModelList(BaseModel):
    """OpenAI model list format."""
    object: str = "list"
    data: list[ModelInfo]


# =============================================================================
# Application State
# =============================================================================

class AppState:
    """Global application state."""
    rag_system: Optional[RAGSystem] = None
    tool_manager: Optional[ToolManager] = None
    llm_client: Optional[AsyncOpenAI] = None
    startup_time: float = time.time()


state = AppState()


# =============================================================================
# Lifespan Management
# =============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifespan - startup and shutdown."""
    log.info("Starting DocRAG server...")

    # Initialize RAG system
    try:
        log.info("Initializing RAG system...")
        state.rag_system = await get_rag_system(
            embedding_model=config.EMBEDDING_MODEL,
            vector_store_path=config.VECTOR_STORE_PATH,
            documents_path=config.DOCUMENTS_PATH,
            chunk_size=config.CHUNK_SIZE,
            chunk_overlap=config.CHUNK_OVERLAP,
        )
        log.info("RAG system initialized successfully")
    except Exception as e:
        log.warning(f"RAG system initialization deferred: {e}")
        state.rag_system = None

    # Initialize tool manager
    try:
        log.info("Initializing tool manager...")
        state.tool_manager = get_tool_manager()
        log.info(f"Tool manager initialized with tools: {state.tool_manager.list_tools()}")
    except Exception as e:
        log.warning(f"Tool manager initialization failed: {e}")
        state.tool_manager = None

    # Initialize OpenRouter client for upstream LLM
    # Debug: Show .env file status
    log.info(f"Looking for .env file at: {ENV_FILE}")
    log.info(f".env file exists: {ENV_FILE.exists()}")
    
    api_key = config.OPENROUTER_API_KEY
    if api_key:
        key_preview = f"{api_key[:8]}...{api_key[-4:]}" if len(api_key) > 12 else "***"
        log.info(f"OPENROUTER_API_KEY found: {key_preview}")
    else:
        log.warning("OPENROUTER_API_KEY not found in environment!")
        log.warning(f"Checked .env at: {ENV_FILE}")
        log.warning("Set OPENROUTER_API_KEY in .env file or as environment variable")
    
    try:
        if api_key:
            log.info("Initializing OpenRouter client...")
            # Create custom httpx client to avoid proxy issues
            import httpx
            http_client = httpx.AsyncClient(
                timeout=httpx.Timeout(60.0, connect=10.0),
            )
            state.llm_client = AsyncOpenAI(
                api_key=api_key,
                base_url=config.OPENROUTER_BASE_URL,
                http_client=http_client,
            )
            log.info(f"OpenRouter client initialized successfully (model: {config.UPSTREAM_MODEL})")
        else:
            log.warning("No OPENROUTER_API_KEY provided - using mock responses")
            state.llm_client = None
    except Exception as e:
        log.error(f"Failed to initialize OpenRouter client: {e}")
        state.llm_client = None

    log.info(f"DocRAG server started on {config.HOST}:{config.PORT}")

    yield

    # Cleanup
    log.info("Shutting down DocRAG server...")
    if state.rag_system:
        await state.rag_system.close()
    log.info("DocRAG server stopped")


# =============================================================================
# FastAPI Application
# =============================================================================

app = FastAPI(
    title="DocRAG API",
    description="OpenAI-compatible RAG server powered by OpenRouter",
    version="1.0.0",
    lifespan=lifespan,
)

# CORS middleware for Open WebUI compatibility
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# =============================================================================
# OpenAI-Compatible Endpoints
# =============================================================================

@app.get("/v1/models")
@app.get("/models")
async def list_models():
    """List available models (OpenAI-compatible)."""
    return ModelList(
        data=[
            ModelInfo(id=config.MODEL_NAME, owned_by="docrag"),
            ModelInfo(id="DocRAG", owned_by="docrag"),
        ]
    )


@app.get("/v1/models/{model_id}")
@app.get("/models/{model_id}")
async def get_model(model_id: str):
    """Get model information (OpenAI-compatible)."""
    if model_id not in [config.MODEL_NAME, "DocRAG"]:
        raise HTTPException(status_code=404, detail="Model not found")
    return ModelInfo(id=model_id, owned_by="docrag")


@app.post("/v1/chat/completions")
@app.post("/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    """Handle chat completions (OpenAI-compatible)."""
    request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"

    try:
        if request.stream:
            return StreamingResponse(
                stream_chat_completion(request, request_id),
                media_type="text/event-stream",
            )
        else:
            return await complete_chat(request, request_id)
    except Exception as e:
        log.exception("Chat completion failed")
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# URL Detection and Website Download
# =============================================================================

def extract_urls_from_message(message: str) -> list[str]:
    """Extract URLs from a message, including domains without scheme."""
    # Match full URLs
    url_pattern = r'https?://[^\s<>"{}|\\^`\[\]]+'
    full_urls = re.findall(url_pattern, message)
    
    # Match domain names (with or without www)
    domain_pattern = r'(?:^|[^\w/-])(?:www\.)?([a-zA-Z0-9][-a-zA-Z0-9]*\.[a-zA-Z]{2,}(?:\.[a-zA-Z]{2,})?)(?:/[^\s]*)?(?=[^\w/-]|$)'
    domains = re.findall(domain_pattern, message)
    
    urls = list(full_urls)
    for domain in domains:
        # Check if it's a valid domain (not just a word)
        if '.' in domain and len(domain) > 4:
            normalized = f"https://{domain}" if not domain.startswith(('http://', 'https://')) else domain
            if normalized not in urls:
                urls.append(normalized)
    
    return urls


def should_download_website(message: str, urls: list[str]) -> bool:
    """Determine if the user wants to access content from a website."""
    if not urls:
        return False
    
    message_lower = message.lower()
    
    # Skip if this is an automated Open WebUI task (not a real user request)
    automated_task_indicators = [
        '### task:', 'generate a concise', 'generate 1-3', 
        'suggest 3-5 relevant follow-up', 'summarizing the chat history',
        'categorizing the main themes', 'follow-up questions or prompts',
    ]
    if any(indicator in message_lower for indicator in automated_task_indicators):
        log.info("Skipping website download - appears to be an automated task")
        return False
    
    # Keywords indicating user wants website content
    access_keywords = [
        'go to', 'visit', 'check', 'look at', 'browse', 'open',
        'what is on', 'tell me about', 'give me', 'show me', 'get',
        'headlines', 'content', 'information from', 'from the', 'from',
        'on the website', 'on the site', 'the website', 'the site', 'website',
        'summarize', 'extract', 'read', 'analyze', 'find', 'search',
        'what does', 'what\'s on', 'what is', 'tell me', 'about',
        'news', 'articles', 'posts', 'page', 'pages',
    ]
    
    # Check if message contains access intent
    has_access_intent = any(kw in message_lower for kw in access_keywords)
    
    # Also trigger if URL is directly mentioned with a question
    has_question = '?' in message or any(qw in message_lower for qw in ['what', 'how', 'who', 'where', 'when', 'why'])
    
    return (has_access_intent or has_question) and len(urls) > 0


async def download_website_if_needed(user_message: str) -> dict[str, Any]:
    """
    Download website if user is asking about one.
    Returns download info if successful.
    """
    urls = extract_urls_from_message(user_message)
    
    if not should_download_website(user_message, urls):
        return {"downloaded": False, "reason": "No website access intent detected"}
    
    if not state.rag_system:
        return {"downloaded": False, "reason": "RAG system not initialized"}
    
    for url in urls:
        try:
            # Check if site is already downloaded
            site_info = state.rag_system.get_site_info(url)
            if site_info:
                log.info(f"Site already downloaded: {url} ({site_info.get('chunks_ingested', 0)} chunks)")
                return {
                    "downloaded": True,
                    "url": url,
                    "chunks": site_info.get("chunks_ingested", 0),
                    "pages": site_info.get("pages_downloaded", 0),
                    "local_path": site_info.get("local_path"),
                    "cached": True,
                }
            
            log.info(f"Auto-downloading website: {url}")
            result = await state.rag_system.download_and_ingest_website(
                url=url,
                max_pages=30,
            )
            
            if result.get("success"):
                log.info(f"Successfully downloaded {url}: {result.get('total_chunks', 0)} chunks")
                return {
                    "downloaded": True,
                    "url": url,
                    "chunks": result.get("total_chunks", 0),
                    "pages": result.get("pages_processed", 0),
                    "local_path": result.get("local_path"),
                }
        except Exception as e:
            log.warning(f"Failed to download {url}: {e}")
    
    return {"downloaded": False, "reason": "All download attempts failed"}


# =============================================================================
# Chat Completion Logic
# =============================================================================

async def complete_chat(request: ChatCompletionRequest, request_id: str) -> ChatCompletionResponse:
    """Process a non-streaming chat completion request."""
    log.info(f"=== Starting complete_chat for request {request_id} ===")
    messages = request.messages

    # Extract the last user message
    user_message = ""
    for msg in reversed(messages):
        if msg.role == "user" and msg.content:
            user_message = msg.content
            break

    if not user_message:
        raise HTTPException(status_code=400, detail="No user message found")

    log.info(f"User message: {user_message[:100]}...")

    # Step 1: Download website if user is asking about one (BEFORE RAG retrieval)
    download_info = await download_website_if_needed(user_message)
    if download_info.get("downloaded"):
        log.info(f"Website auto-downloaded: {download_info.get('url')}")

    # Step 2: RAG Retrieval (now includes newly downloaded content)
    context = ""
    sources = []
    if state.rag_system:
        try:
            rag_result = await state.rag_system.query(
                query=user_message,
                top_k=config.TOP_K_RESULTS,
            )
            context = rag_result.get("context", "")
            sources = rag_result.get("sources", [])
            log.info(f"RAG retrieved {len(sources)} relevant documents")
        except Exception as e:
            log.warning(f"RAG retrieval failed: {e}")

    # Step 3: Build enhanced prompt with context
    enhanced_messages = build_enhanced_messages(messages, context, sources, download_info)

    # Step 4: Generate response with upstream LLM
    log.info(f"Calling generate_response for request {request_id}")
    response_content = await generate_response(
        enhanced_messages,
        temperature=request.temperature,
        max_tokens=request.max_tokens,
    )
    log.info(f"=== Completed complete_chat for request {request_id} ===")

    # Step 5: Build and return response
    return ChatCompletionResponse(
        id=request_id,
        model=config.MODEL_NAME,
        choices=[
            ChatCompletionChoice(
                message=ChatMessage(
                    role="assistant",
                    content=response_content,
                ),
                finish_reason="stop",
            )
        ],
        usage=ChatCompletionUsage(
            prompt_tokens=len(str(enhanced_messages)) // 4,
            completion_tokens=len(response_content) // 4,
        ),
    )


async def stream_chat_completion(
    request: ChatCompletionRequest,
    request_id: str,
) -> AsyncIterator[str]:
    """Stream a chat completion response."""
    messages = request.messages

    # Extract the last user message
    user_message = ""
    for msg in reversed(messages):
        if msg.role == "user" and msg.content:
            user_message = msg.content
            break

    if not user_message:
        yield f"data: {json.dumps({'error': 'No user message found'})}\n\n"
        return

    # Step 1: Download website if user is asking about one (BEFORE RAG retrieval)
    download_info = await download_website_if_needed(user_message)
    if download_info.get("downloaded"):
        log.info(f"Website auto-downloaded: {download_info.get('url')}")

    # Step 2: RAG Retrieval (now includes newly downloaded content)
    context = ""
    sources = []
    if state.rag_system:
        try:
            rag_result = await state.rag_system.query(
                query=user_message,
                top_k=config.TOP_K_RESULTS,
            )
            context = rag_result.get("context", "")
            sources = rag_result.get("sources", [])
            log.info(f"RAG retrieved {len(sources)} relevant documents")
        except Exception as e:
            log.warning(f"RAG retrieval failed: {e}")

    # Step 3: Build enhanced prompt with context
    enhanced_messages = build_enhanced_messages(messages, context, sources, download_info)

    # Step 4: Stream response from upstream LLM
    created = int(time.time())

    try:
        if state.llm_client:
            # For streaming with tools, we need to handle tool calls first
            # Then stream the final response
            if state.tool_manager and config.ENABLE_TOOLS:
                # Use non-streaming for tool calls, then stream the result
                response_content = await generate_response(
                    enhanced_messages,
                    temperature=request.temperature or 0.7,
                    max_tokens=request.max_tokens or 4096,
                )
                # Stream the final response as a single chunk
                yield f"data: {json.dumps({
                    'id': request_id,
                    'object': 'chat.completion.chunk',
                    'created': created,
                    'model': config.MODEL_NAME,
                    'choices': [{
                        'index': 0,
                        'delta': {'content': response_content},
                        'finish_reason': None
                    }]
                })}\n\n"
            else:
                # No tools - use regular streaming
                stream = await state.llm_client.chat.completions.create(
                    model=config.UPSTREAM_MODEL,
                    messages=[{"role": m.role, "content": m.content} for m in enhanced_messages if m.content],
                    temperature=request.temperature or 0.7,
                    max_tokens=request.max_tokens or 4096,
                    stream=True,
                )

                async for chunk in stream:
                    if chunk.choices and chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        yield f"data: {json.dumps({
                            'id': request_id,
                            'object': 'chat.completion.chunk',
                            'created': created,
                            'model': config.MODEL_NAME,
                            'choices': [{
                                'index': 0,
                                'delta': {'content': content},
                                'finish_reason': None
                            }]
                        })}\n\n"

            # Send final chunk
            yield f"data: {json.dumps({
                'id': request_id,
                'object': 'chat.completion.chunk',
                'created': created,
                'model': config.MODEL_NAME,
                'choices': [{
                    'index': 0,
                    'delta': {},
                    'finish_reason': 'stop'
                }]
            })}\n\n"
            yield "data: [DONE]\n\n"

        else:
            # Mock streaming response for testing
            mock_response = f"I understand you're asking about: {user_message}\n\n"
            if download_info.get("downloaded"):
                mock_response += f"I have downloaded and analyzed {download_info.get('url')}.\n"
                mock_response += f"Processed {download_info.get('pages')} pages into {download_info.get('chunks')} chunks.\n\n"
            if context:
                mock_response += f"Based on my knowledge base, here's what I found:\n\n{context[:1000]}...\n\n"
            mock_response += "\n\n[Demo mode - configure OPENROUTER_API_KEY for full LLM responses]"

            for char in mock_response:
                yield f"data: {json.dumps({
                    'id': request_id,
                    'object': 'chat.completion.chunk',
                    'created': created,
                    'model': config.MODEL_NAME,
                    'choices': [{
                        'index': 0,
                        'delta': {'content': char},
                        'finish_reason': None
                    }]
                })}\n\n"
                await asyncio.sleep(0.01)

            yield f"data: {json.dumps({
                'id': request_id,
                'object': 'chat.completion.chunk',
                'created': created,
                'model': config.MODEL_NAME,
                'choices': [{
                    'index': 0,
                    'delta': {},
                    'finish_reason': 'stop'
                }]
            })}\n\n"
            yield "data: [DONE]\n\n"

    except Exception as e:
        log.exception("Streaming failed")
        yield f"data: {json.dumps({'error': str(e)})}\n\n"


def build_enhanced_messages(
    messages: list[ChatMessage],
    context: str,
    sources: list[str],
    download_info: dict = None,
    tool_results: list[dict] = None,
) -> list[ChatMessage]:
    """Build enhanced messages with RAG context."""
    enhanced = []

    # Build tool descriptions for context
    tool_descriptions = _build_tool_descriptions()
    
    # Add system message with RAG context and tool instructions
    system_content = f"""You are a helpful AI assistant with access to real-time data through various tools.

## AVAILABLE TOOLS
{tool_descriptions}

## HOW TO USE TOOLS
When you need to use one or more tools, output a SINGLE JSON block containing ALL tool calls as an array.
You MUST bundle every tool call into one response - do NOT respond with just one tool at a time.

Output EXACTLY this format (nothing else before or after):
```json
{{"tool_calls": [
  {{"name": "tool_name", "arguments": {{"arg1": "value1"}}}},
  {{"name": "another_tool", "arguments": {{"arg2": "value2"}}}}
]}}
```

## IMPORTANT RULES
1. ALWAYS use tools to get CURRENT data - do NOT say you cannot access real-time data
2. When asked about stocks, crypto, weather, or news, you MUST use the appropriate tool(s)
3. Bundle ALL needed tool calls into a single `tool_calls` array - include every tool you need in one response
4. After receiving tool results, provide a helpful, natural-language response based on the data
5. Be concise and factual - report exact data from tools
"""

    if download_info and download_info.get("downloaded"):
        system_content += f"\n\n--- Website Access ---\n"
        system_content += f"Downloaded website: {download_info.get('url')}\n"
        system_content += f"Pages: {download_info.get('pages')}, Chunks: {download_info.get('chunks')}\n"

    if context:
        system_content += f"\n\n--- Relevant Context from Knowledge Base ---\n{context}\n"
        if sources:
            system_content += f"\n--- Sources ---\n" + "\n".join(f"- {s}" for s in sources[:10])

    # Add previous tool results if any
    if tool_results:
        system_content += "\n\n--- PREVIOUS TOOL RESULTS ---\n"
        for tr in tool_results:
            system_content += f"\nTool: {tr['name']}\nResult: {json.dumps(tr['result'], indent=2)}\n"

    enhanced.append(ChatMessage(role="system", content=system_content))

    # Add conversation history (excluding old system messages)
    for msg in messages:
        if msg.role != "system":
            enhanced.append(msg)

    return enhanced


def _build_tool_descriptions() -> str:
    """Build a concise description of all available tools for the system prompt."""
    if not state.tool_manager:
        return "No tools available."
    
    descriptions = []
    for name, schema in state.tool_manager._schemas.items():
        func = schema.get("function", {})
        desc = func.get("description", "")[:100]  # Truncate long descriptions
        params = func.get("parameters", {}).get("properties", {})
        required = func.get("parameters", {}).get("required", [])
        
        # Build param list
        param_strs = []
        for pname, pinfo in params.items():
            ptype = pinfo.get("type", "any")
            preq = " (required)" if pname in required else ""
            param_strs.append(f"{pname}: {ptype}{preq}")
        
        params_str = ", ".join(param_strs) if param_strs else "none"
        descriptions.append(f"- {name}({params_str}): {desc}")
    
    return "\n".join(descriptions)


def _parse_tool_calls(content: str) -> list[dict]:
    """Parse tool calls from LLM response content.

    Expects the LLM to output a JSON block like:
    ```json
    {"tool_calls": [{"name": "tool_name", "arguments": {...}}, ...]}
    ```

    Returns a list of tool call dicts, each with 'name' and 'arguments' keys.
    """
    tool_calls = []

    def _extract_json_object(text: str, start_key: str) -> Optional[dict]:
        """Extract a JSON object containing start_key using brace counting."""
        idx = text.find(start_key)
        if idx == -1:
            return None
        # Walk backwards to find the opening {
        depth = 0
        obj_start = -1
        for i in range(idx, -1, -1):
            if text[i] == '}':
                depth += 1
            elif text[i] == '{':
                if depth == 0:
                    obj_start = i
                    break
                depth -= 1
        if obj_start == -1:
            return None
        # Walk forwards to find the matching closing }
        depth = 0
        obj_end = -1
        for i in range(obj_start, len(text)):
            if text[i] == '{':
                depth += 1
            elif text[i] == '}':
                depth -= 1
                if depth == 0:
                    obj_end = i + 1
                    break
        if obj_end == -1:
            return None
        try:
            return json.loads(text[obj_start:obj_end])
        except json.JSONDecodeError:
            return None

    # --- Pattern 1: {"tool_calls": [...]} in a code fence block ---
    fence_matches = re.findall(r'```\w*\s*(.*?)\s*```', content, re.DOTALL)
    for block_text in fence_matches:
        obj = _extract_json_object(block_text, '"tool_calls"')
        if obj and "tool_calls" in obj and isinstance(obj["tool_calls"], list):
            for tc in obj["tool_calls"]:
                if isinstance(tc, dict) and "name" in tc:
                    tool_calls.append(tc)
            if tool_calls:
                return tool_calls

    # --- Pattern 2: {"tool_calls": [...]} bare JSON (outside code fences) ---
    stripped = re.sub(r'```\w*\s*.*?\s*```', '', content, flags=re.DOTALL)
    obj = _extract_json_object(stripped, '"tool_calls"')
    if obj and "tool_calls" in obj and isinstance(obj["tool_calls"], list):
        for tc in obj["tool_calls"]:
            if isinstance(tc, dict) and "name" in tc:
                tool_calls.append(tc)
        if tool_calls:
            return tool_calls

    # --- Pattern 3 (legacy fallback): {"tool_call": {...}} single tool ---
    # Also support the old format in case the LLM ignores instructions
    for block_text in fence_matches:
        obj = _extract_json_object(block_text, '"tool_call"')
        if obj and "tool_call" in obj and isinstance(obj["tool_call"], dict) and "name" in obj["tool_call"]:
            tool_calls.append(obj["tool_call"])
    if not tool_calls:
        obj = _extract_json_object(stripped, '"tool_call"')
        if obj and "tool_call" in obj and isinstance(obj["tool_call"], dict) and "name" in obj["tool_call"]:
            tool_calls.append(obj["tool_call"])

    return tool_calls


async def generate_response(
    messages: list[ChatMessage],
    temperature: float = 0.7,
    max_tokens: int = 4096,
) -> str:
    """Generate response using upstream LLM via OpenRouter.

    Uses content-based tool calling: the LLM outputs a single JSON block with
    all tool calls bundled as a `tool_calls` array. This works around model
    limitations on the number of native tool calls per response.
    """
    if not state.llm_client:
        # Mock response for testing
        user_msg = ""
        for msg in reversed(messages):
            if msg.role == "user" and msg.content:
                user_msg = msg.content
                break
        return f"Demo mode response. Your question: {user_msg[:100]}... Configure OPENROUTER_API_KEY for full functionality."

    try:
        # Convert messages to dict format
        messages_dict = []
        for m in messages:
            if m.content:
                messages_dict.append({"role": m.role, "content": m.content})

        # Tool calling loop (content-based approach — no `tools` param to API)
        max_iterations = config.MAX_TOOL_ITERATIONS
        iteration = 0

        while iteration < max_iterations:
            iteration += 1
            log.info(f"LLM call iteration {iteration}")

            # Call LLM WITHOUT tools parameter — tool instructions are in the system prompt
            response = await state.llm_client.chat.completions.create(
                model=config.UPSTREAM_MODEL,
                messages=messages_dict,
                temperature=temperature,
                max_tokens=max_tokens,
            )

            if not response.choices:
                log.warning("No choices in response")
                return "I apologize, but I couldn't generate a response."

            content = response.choices[0].message.content or ""
            log.info(f"LLM response: content_len={len(content)}")

            # --- Parse tool calls from content ---
            tool_calls = _parse_tool_calls(content)

            if tool_calls:
                log.info(f"Parsed {len(tool_calls)} tool calls from content")

                # Execute ALL tools concurrently
                if state.tool_manager:
                    import asyncio as _asyncio

                    async def _run_tool(tc):
                        name = tc.get("name")
                        args = tc.get("arguments", {})
                        if not isinstance(args, dict):
                            try:
                                args = json.loads(args)
                            except (json.JSONDecodeError, TypeError):
                                args = {}
                        result = await _asyncio.to_thread(
                            state.tool_manager.execute_tool, name, args
                        )
                        return name, result

                    results = await _asyncio.gather(*[_run_tool(tc) for tc in tool_calls])

                    # Build a single consolidated results block
                    results_text = ""
                    for name, result in results:
                        log.info(f"Tool {name} result: success={result.get('success', False)}")
                        results_text += f"\n### Tool: {name}\n{json.dumps(result, indent=2)}\n"

                    # Append assistant's tool call message to conversation
                    messages_dict.append({"role": "assistant", "content": content})

                    # Feed ALL results back in one user message
                    messages_dict.append({
                        "role": "user",
                        "content": (
                            f"--- ALL TOOL RESULTS ---\n"
                            f"Executed {len(tool_calls)} tool(s). Results:\n{results_text}\n"
                            f"---\n\n"
                            f"Now provide a helpful response to the original question using ALL the data above."
                        ),
                    })

                    continue
                else:
                    log.warning("Tool call detected but tool_manager is None")

            # --- No tool calls — return the final response ---
            cleaned_content = _clean_tool_syntax(content)
            log.info(f"Returning final response (len={len(cleaned_content)})")
            return cleaned_content or "I apologize, but I couldn't generate a response."

        # Max iterations reached
        log.warning(f"Max iterations ({max_iterations}) reached")
        return "I reached the maximum number of tool call rounds. Please try a more specific question."

    except Exception as e:
        log.error(f"OpenRouter LLM call failed: {e}")
        import traceback
        log.error(traceback.format_exc())
        return f"I encountered an error: {str(e)}"


def _clean_tool_syntax(content: str) -> str:
    """Remove tool call JSON blocks from response text.

    Strips code-fence-wrapped blocks containing "tool_calls" or "tool_call".
    Does NOT strip bare JSON to avoid removing valid content.
    """
    def remove_code_block(m):
        block = m.group(0)
        inner = m.group(1)
        if '"tool_calls"' in inner or '"tool_call"' in inner:
            return ''
        return block

    cleaned = re.sub(r'```\w*\s*(.*?)\s*```', remove_code_block, content, flags=re.DOTALL)
    return cleaned.strip()


# =============================================================================
# Document Management Endpoints
# =============================================================================

@app.post("/v1/documents/upload")
async def upload_document(request: Request):
    """Upload a document to the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        form = await request.form()
        file = form.get("file")
        if not file:
            raise HTTPException(status_code=400, detail="No file provided")

        content = await file.read()
        filename = file.filename or "unknown"

        # Process and store document
        result = await state.rag_system.add_document(
            content=content,
            filename=filename,
        )

        return {"success": True, "message": f"Document '{filename}' added", "chunks": result.get("chunks", 0)}

    except Exception as e:
        log.exception("Document upload failed")
        raise HTTPException(status_code=500, detail=str(e))


class WebsiteDownloadRequest(BaseModel):
    """Request model for website download."""
    url: str
    max_pages: int = 50
    threads: int = 6
    download_external_assets: bool = False
    external_domains: Optional[list[str]] = None


@app.post("/v1/documents/website")
async def download_website(request: WebsiteDownloadRequest):
    """
    Download a website and ingest it into the knowledge base.
    
    This is the PRIMARY way to add content to the RAG system.
    Uses the website_downloader_tool to download and process websites.
    """
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        log.info(f"Downloading website: {request.url}")
        
        result = await state.rag_system.download_and_ingest_website(
            url=request.url,
            max_pages=request.max_pages,
            threads=request.threads,
            download_external_assets=request.download_external_assets,
            external_domains=request.external_domains,
        )

        if result.get("success"):
            return {
                "success": True,
                "message": f"Website downloaded and ingested: {request.url}",
                "url": request.url,
                "local_path": result.get("local_path"),
                "pages_processed": result.get("pages_processed", 0),
                "total_chunks": result.get("total_chunks", 0),
            }
        else:
            raise HTTPException(
                status_code=500,
                detail=result.get("message", "Website download failed")
            )

    except HTTPException:
        raise
    except Exception as e:
        log.exception("Website download failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/v1/documents/url")
async def add_document_from_url(request: dict):
    """Add a document from URL to the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    url = request.get("url")
    if not url:
        raise HTTPException(status_code=400, detail="No URL provided")

    try:
        result = await state.rag_system.add_document_from_url(url)
        return {"success": True, "message": f"Document from {url} added", "chunks": result.get("chunks", 0)}
    except Exception as e:
        log.exception("URL document addition failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/documents")
async def list_documents():
    """List documents in the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        docs = await state.rag_system.list_documents()
        return {"documents": docs}
    except Exception as e:
        log.exception("Document listing failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/documents/sites")
async def list_downloaded_sites():
    """List all downloaded websites in the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        sites = await state.rag_system.list_downloaded_sites()
        return {"sites": sites}
    except Exception as e:
        log.exception("Site listing failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/documents/sites/{url:path}")
async def get_site_info(url: str):
    """Get information about a specific downloaded site."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        from urllib.parse import unquote
        decoded_url = unquote(url)
        
        if not decoded_url.startswith(("http://", "https://")):
            decoded_url = "https://" + decoded_url
        
        site_info = state.rag_system.get_site_info(decoded_url)
        if site_info:
            return {"site": site_info}
        else:
            raise HTTPException(status_code=404, detail="Site not found")
    except HTTPException:
        raise
    except Exception as e:
        log.exception("Site info retrieval failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/v1/documents/{doc_id}")
async def delete_document(doc_id: str):
    """Delete a document from the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        await state.rag_system.delete_document(doc_id)
        return {"success": True, "message": f"Document {doc_id} deleted"}
    except Exception as e:
        log.exception("Document deletion failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/v1/documents/sites/{url:path}")
async def delete_site(url: str):
    """Delete a downloaded website and all its content from the knowledge base."""
    if not state.rag_system:
        raise HTTPException(status_code=503, detail="RAG system not initialized")

    try:
        from urllib.parse import unquote
        decoded_url = unquote(url)
        
        if not decoded_url.startswith(("http://", "https://")):
            decoded_url = "https://" + decoded_url
        
        result = await state.rag_system.delete_site(decoded_url)
        
        if result.get("success"):
            return {
                "success": True,
                "message": f"Site {decoded_url} deleted",
                "deleted_chunks": result.get("deleted_chunks", 0),
                "deleted_path": result.get("deleted_path"),
            }
        else:
            raise HTTPException(status_code=404, detail=result.get("message", "Site not found"))
            
    except HTTPException:
        raise
    except Exception as e:
        log.exception("Site deletion failed")
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Health and Status Endpoints
# =============================================================================

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "uptime": time.time() - state.startup_time,
        "rag_enabled": state.rag_system is not None,
        "tools_enabled": state.tool_manager is not None,
        "llm_connected": state.llm_client is not None,
    }


@app.get("/")
async def root():
    """Root endpoint with API info."""
    return {
        "name": "DocRAG API",
        "version": "1.0.0",
        "description": "OpenAI-compatible RAG server powered by OpenRouter. Auto-downloads and analyzes websites when users ask about them.",
        "endpoints": {
            "chat": "/v1/chat/completions",
            "models": "/v1/models",
            "documents": "/v1/documents",
            "download_website": "/v1/documents/website",
            "list_sites": "/v1/documents/sites",
            "health": "/health",
        },
    }


# =============================================================================
# Main Entry Point
# =============================================================================

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "main:app",
        host=config.HOST,
        port=config.PORT,
        reload=config.DEBUG,
    )