app.py

import os
import json
import asyncio
import requests
from datetime import datetime
from typing import List, Dict, Optional
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from openai import OpenAI
import logging

# --- Security Helper Functions ---
def verify_origin(request: Request):
    """Verify that the request comes from an allowed origin for /chat endpoint"""
    origin = request.headers.get("origin")
    referer = request.headers.get("referer")
    
    allowed_origins = [
        "https://chrunos.com",
        "https://www.chrunos.com"
    ]
    
    # Allow localhost for development (you can remove this in production)
    if origin and any(origin.startswith(local) for local in ["http://localhost:", "http://127.0.0.1:"]):
        return True
    
    # Check origin header
    if origin in allowed_origins:
        return True
    
    # Check referer header as fallback
    if referer and any(referer.startswith(allowed) for allowed in allowed_origins):
        return True
    
    raise HTTPException(
        status_code=403, 
        detail="Access denied: This endpoint is only accessible from chrunos.com"
    )

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Load API Keys from Environment Variables ---
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
GOOGLE_CX = os.getenv("GOOGLE_CX")
LLM_API_KEY = os.getenv("LLM_API_KEY")
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api-15i2e8ze256bvfn6.aistudio-app.com/v1")

# --- Improved System Prompts ---
SYSTEM_PROMPT_WITH_SEARCH = """You are an intelligent AI assistant with access to real-time web search capabilities.

When you need current information, recent events, specific facts, or when the user's question would benefit from up-to-date information, use the google_search function.

**Use search for:**
- Recent news or events
- Current statistics or data
- Specific factual information you're unsure about
- Questions about things that may have changed recently
- When the user explicitly asks for current/recent information

**Response Guidelines:**
1. Always use the search tool when it would provide more accurate or current information
2. Synthesize information from multiple sources when available
3. Clearly indicate when information comes from search results
4. Provide comprehensive, well-structured answers
5. Cite sources appropriately with links.
6. If search results conflict with my knowledge, prioritize the search results.

Current date: {current_date}"""

SYSTEM_PROMPT_NO_SEARCH = """You are an intelligent AI assistant. Provide helpful, accurate, and comprehensive responses based on your training data.

Current date: {current_date}"""

# --- Optimized Web Search Tool ---
async def google_search_tool_async(query: str, num_results: int = 3) -> List[Dict]:
    """
    Async Google Custom Search - reduced results for faster response
    """
    if not GOOGLE_API_KEY or not GOOGLE_CX or not query.strip():
        return []

    logger.info(f"Executing search for: '{query}'")
    
    search_url = "https://www.googleapis.com/customsearch/v1"
    params = {
        "key": GOOGLE_API_KEY, 
        "cx": GOOGLE_CX, 
        "q": query.strip(), 
        "num": min(num_results, 5),
        "dateRestrict": "m3"
    }
    
    try:
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None, 
            lambda: requests.get(search_url, params=params, timeout=10)
        )
        response.raise_for_status()
        search_results = response.json()
        
        if "items" not in search_results:
            return []
        
        parsed_results = []
        for item in search_results.get("items", [])[:num_results]:
            title = item.get("title", "").strip()
            url = item.get("link", "").strip()
            snippet = item.get("snippet", "").strip()
            
            if title and url and snippet:
                parsed_results.append({
                    "source_title": title,
                    "url": url,
                    "snippet": snippet,
                    "domain": url.split('/')[2] if '/' in url else url
                })
            
        logger.info(f"Retrieved {len(parsed_results)} search results")
        return parsed_results
        
    except Exception as e:
        logger.error(f"Search error: {e}")
        return []

def format_search_results_compact(search_results: List[Dict]) -> str:
    """Compact formatting for faster processing"""
    if not search_results:
        return "No search results found."
    
    formatted = ["Search Results:"]
    for i, result in enumerate(search_results, 1):
        formatted.append(f"\n{i}. {result['source_title']}")
        formatted.append(f"   Source: {result['domain']}")
        formatted.append(f"   Content: {result['snippet']}")
    
    return "\n".join(formatted)

# --- FastAPI Application Setup ---
app = FastAPI(title="Streaming AI Chatbot", version="2.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://chrunos.com",
        "https://www.chrunos.com",
        "http://localhost:3000",
        "http://localhost:8000",
    ],
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

# --- OpenAI Client Initialization ---
if not LLM_API_KEY or not LLM_BASE_URL:
    logger.error("LLM_API_KEY or LLM_BASE_URL not configured")
    client = None
else:
    client = OpenAI(api_key=LLM_API_KEY, base_url=LLM_BASE_URL)
    logger.info("OpenAI client initialized successfully")

# --- Tool Definition ---
available_tools = [
    {
        "type": "function", 
        "function": {
            "name": "google_search", 
            "description": "Search Google for current information, recent events, or specific facts. Use this when you need up-to-date information or when the user's question would benefit from current data.",
            "parameters": {
                "type": "object", 
                "properties": {
                    "query": {
                        "type": "string", 
                        "description": "Search query with relevant keywords"
                    }
                }, 
                "required": ["query"]
            }
        }
    }
]

# --- Fixed Streaming Response Generator ---
async def generate_streaming_response(messages: List[Dict], use_search: bool, temperature: float):
    """Generate streaming response with optional search"""
    
    try:
        # Initial LLM call with streaming
        llm_kwargs = {
            "model": "unsloth/Qwen3-30B-A3B-GGUF", 
            "temperature": temperature,
            "messages": messages,
            "max_tokens": 2000,
            "stream": True
        }
        
        if use_search:
            llm_kwargs["tools"] = available_tools
            llm_kwargs["tool_choice"] = "auto"
        
        source_links = []
        response_content = ""
        tool_calls_data = []
        current_tool_call = None
        
        # First streaming call
        stream = client.chat.completions.create(**llm_kwargs)
        
        # Track if we're in the middle of collecting a tool call
        collecting_tool_call = False
        
        for chunk in stream:
            delta = chunk.choices[0].delta
            finish_reason = chunk.choices[0].finish_reason
            
            # Handle content streaming
            if delta.content:
                content_chunk = delta.content
                response_content += content_chunk
                yield f"data: {json.dumps({'type': 'content', 'data': content_chunk})}\n\n"
            
            # Handle tool calls - FIXED LOGIC
            if delta.tool_calls:
                collecting_tool_call = True
                for tool_call in delta.tool_calls:
                    # Ensure we have enough slots in tool_calls_data
                    while len(tool_calls_data) <= tool_call.index:
                        tool_calls_data.append({
                            "id": None,
                            "function": {"name": None, "arguments": ""}
                        })
                    
                    # Update the tool call data
                    if tool_call.id:
                        tool_calls_data[tool_call.index]["id"] = tool_call.id
                    if tool_call.function and tool_call.function.name:
                        tool_calls_data[tool_call.index]["function"]["name"] = tool_call.function.name
                    if tool_call.function and tool_call.function.arguments:
                        tool_calls_data[tool_call.index]["function"]["arguments"] += tool_call.function.arguments
            
            # Check if we've finished collecting tool calls
            if finish_reason in ["tool_calls", "stop"] and collecting_tool_call:
                break
        
        # Process tool calls if any were collected
        processed_any_tools = False
        if tool_calls_data and any(tc.get("id") and tc.get("function", {}).get("name") for tc in tool_calls_data):
            yield f"data: {json.dumps({'type': 'status', 'data': 'Searching...'})}\n\n"
            
            tool_responses = []
            
            # Process each tool call
            for tool_call in tool_calls_data:
                if not tool_call.get("id") or not tool_call.get("function", {}).get("name"):
                    continue
                    
                function_name = tool_call["function"]["name"]
                
                if function_name == "google_search":
                    try:
                        args = json.loads(tool_call["function"]["arguments"])
                        query = args.get("query", "").strip()
                        if query:
                            logger.info(f"Executing search with query: {query}")
                            search_results = await google_search_tool_async(query)
                            
                            if search_results:
                                processed_any_tools = True
                                
                                # Collect source links
                                for result in search_results:
                                    source_links.append({
                                        "title": result["source_title"],
                                        "url": result["url"],
                                        "domain": result["domain"]
                                    })
                                
                                # Format results for the model
                                search_context = format_search_results_compact(search_results)
                                tool_responses.append({
                                    "tool_call_id": tool_call["id"],
                                    "role": "tool",
                                    "content": search_context
                                })
                            else:
                                tool_responses.append({
                                    "tool_call_id": tool_call["id"],
                                    "role": "tool",
                                    "content": "No search results found."
                                })
                    except json.JSONDecodeError as e:
                        logger.error(f"Failed to parse tool arguments: {e}")
                        tool_responses.append({
                            "tool_call_id": tool_call["id"],
                            "role": "tool",
                            "content": "Error: Invalid search query format."
                        })
                    except Exception as e:
                        logger.error(f"Search tool error: {e}")
                        tool_responses.append({
                            "tool_call_id": tool_call["id"],
                            "role": "tool",
                            "content": f"Search error: {str(e)}"
                        })
            
            # If we have tool responses, make a second call to get the final response
            if tool_responses:
                yield f"data: {json.dumps({'type': 'status', 'data': 'Generating response...'})}\n\n"
                
                # Add tool call and tool response messages
                final_messages = messages.copy()
                
                # Add the assistant's tool call message
                assistant_message = {
                    "role": "assistant",
                    "content": response_content if response_content else None,
                    "tool_calls": [
                        {
                            "id": tc["id"],
                            "type": "function",
                            "function": {
                                "name": tc["function"]["name"],
                                "arguments": tc["function"]["arguments"]
                            }
                        }
                        for tc in tool_calls_data if tc.get("id") and tc.get("function", {}).get("name")
                    ]
                }
                final_messages.append(assistant_message)
                
                # Add tool response messages
                final_messages.extend(tool_responses)
                
                # Generate final response
                final_stream = client.chat.completions.create(
                    model="unsloth/Qwen3-30B-A3B-GGUF",
                    temperature=temperature,
                    messages=final_messages,
                    max_tokens=2000,
                    stream=True
                )
                
                for chunk in final_stream:
                    if chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        yield f"data: {json.dumps({'type': 'content', 'data': content})}\n\n"
        
        # Send sources and completion
        if source_links:
            yield f"data: {json.dumps({'type': 'sources', 'data': source_links})}\n\n"
        
        yield f"data: {json.dumps({'type': 'done', 'data': {'search_used': processed_any_tools}})}\n\n"
        
    except Exception as e:
        logger.error(f"Streaming error: {e}")
        yield f"data: {json.dumps({'type': 'error', 'data': str(e)})}\n\n"

# --- Streaming Chat Endpoint ---
@app.post("/chat/stream")
async def chat_stream_endpoint(request: Request, _: None = Depends(verify_origin)):
    if not client:
        raise HTTPException(status_code=500, detail="LLM client not configured")
        
    try:
        data = await request.json()
        user_message = data.get("message", "").strip()
        use_search = data.get("use_search", False)
        temperature = max(0, min(2, data.get("temperature", 0.7)))
        conversation_history = data.get("history", [])
        user_prompt = data.get("system_prompt")
        
        if not user_message:
            raise HTTPException(status_code=400, detail="No message provided")

        # Prepare messages
        current_date = datetime.now().strftime("%Y-%m-%d")
        system_content = (SYSTEM_PROMPT_WITH_SEARCH if use_search else user_prompt#SYSTEM_PROMPT_NO_SEARCH
                         ).format(current_date=current_date)
        messages = [{"role": "system", "content": system_content}] + conversation_history + [{"role": "user", "content": user_message}]
        
        logger.info(f"Stream request - search: {use_search}, temp: {temperature}, message: {user_message[:100]}...")
        
        return StreamingResponse(
            generate_streaming_response(messages, use_search, temperature),
            media_type="text/plain",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"
            }
        )
        
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON")
    except Exception as e:
        logger.error(f"Stream endpoint error: {e}")
        raise HTTPException(status_code=500, detail=str(e))