Update app.py

app.py

@@ -3,12 +3,38 @@ import json
 import requests
 from datetime import datetime
 from typing import List, Dict, Optional
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request, HTTPException, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from openai import OpenAI
 import logging
-from fastapi.responses import StreamingResponse
 
+# --- 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)
@@ -157,9 +183,14 @@ app = FastAPI(title="AI Chatbot with Enhanced Search", version="2.0.0")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production, specify actual origins
+    allow_origins=[
+        "https://chrunos.com",
+        "https://www.chrunos.com",
+        "http://localhost:3000",  # For local development
+        "http://localhost:8000",  # For local development
+    ],
     allow_credentials=True,
-    allow_methods=["*"],
+    allow_methods=["GET", "POST", "OPTIONS"],
     allow_headers=["*"],
 )
 
@@ -221,51 +252,154 @@ def should_use_search(message: str) -> bool:
 
 # --- Enhanced Chatbot Endpoint ---
 @app.post("/chat")
-async def chat_endpoint(request: Request):
+async def chat_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", data.get("user_search"))
+        
+        # Support both 'use_search' and 'user_search' parameter names for flexibility
+        use_search = data.get("use_search")
+        if use_search is None:
+            use_search = data.get("user_search")  # Alternative parameter name
+        
+        # Allow client to specify temperature (with validation)
+        temperature = data.get("temperature", 0.7)  # Default to 0.7
+        if not isinstance(temperature, (int, float)) or temperature < 0 or temperature > 2:
+            logger.warning(f"Invalid temperature value: {temperature}, defaulting to 0.7")
+            temperature = 0.7
+        
         conversation_history = data.get("history", [])
-
+        
+        # Debug logging for request parameters
+        logger.info(f"Request parameters - message length: {len(user_message)}, use_search: {use_search}, temperature: {temperature}, history length: {len(conversation_history)}")
+        
         if not user_message:
             raise HTTPException(status_code=400, detail="No message provided")
 
+        # Auto-decide search usage if not specified
         if use_search is None:
             use_search = should_use_search(user_message)
+            logger.info(f"Auto-decided search usage: {use_search}")
+        else:
+            logger.info(f"Manual search setting: {use_search}")
 
+        # Prepare messages with appropriate system prompt based on search availability
         current_date = datetime.now().strftime("%Y-%m-%d")
+        
         if use_search:
             system_content = SYSTEM_PROMPT_WITH_SEARCH.format(current_date=current_date)
         else:
             system_content = SYSTEM_PROMPT_NO_SEARCH.format(current_date=current_date)
+            
         system_message = {"role": "system", "content": system_content}
         messages = [system_message] + conversation_history + [{"role": "user", "content": user_message}]
+        
         llm_kwargs = {
-            "model": "unsloth/Qwen3-30B-A3B-GGUF",
-            "temperature": 0.7,
+            "model": "unsloth/Qwen3-30B-A3B-GGUF", 
+            "temperature": temperature,  # Use client-specified temperature
             "messages": messages,
-            "max_tokens": 2000,
-            "stream": True,   # <--- Enable streaming
+            "max_tokens": 2000  # Ensure comprehensive responses
         }
+        
         if use_search:
+            logger.info("Search is ENABLED - tools will be available to the model")
             llm_kwargs["tools"] = available_tools
-            llm_kwargs["tool_choice"] = "auto"
+            llm_kwargs["tool_choice"] = "auto"  # Consider using "required" for testing
+        else:
+            logger.info("Search is DISABLED - no tools available")
+
+        # First LLM call
+        logger.info(f"Making LLM request with tools: {bool(use_search)}, temperature: {temperature}")
+        llm_response = client.chat.completions.create(**llm_kwargs)
+        tool_calls = llm_response.choices[0].message.tool_calls
+        source_links = []
+
+        # Debug: Log tool call information
+        if tool_calls:
+            logger.info(f"LLM made {len(tool_calls)} tool calls")
+            for i, call in enumerate(tool_calls):
+                logger.info(f"Tool call {i+1}: {call.function.name} with args: {call.function.arguments}")
+        else:
+            logger.info("LLM did not make any tool calls")
+            if use_search:
+                logger.warning("Search was enabled but LLM chose not to use search tools - this might indicate the query doesn't require current information")
+
+        if tool_calls:
+            logger.info(f"Processing {len(tool_calls)} tool calls")
+            tool_outputs = []
+            
+            for tool_call in tool_calls:
+                if tool_call.function.name == "google_search":
+                    try:
+                        function_args = json.loads(tool_call.function.arguments)
+                        search_query = function_args.get("query", "").strip()
+                        
+                        if search_query:
+                            logger.info(f"Executing search for: {search_query}")
+                            search_results = google_search_tool([search_query], num_results=5)
+                            
+                            # Collect source links for response
+                            for result in search_results:
+                                source_links.append({
+                                    "title": result["source_title"],
+                                    "url": result["url"],
+                                    "domain": result["domain"]
+                                })
+                            
+                            # Format results for LLM
+                            formatted_results = format_search_results_for_llm(search_results)
+                            tool_outputs.append({
+                                "tool_call_id": tool_call.id, 
+                                "output": formatted_results
+                            })
+                        else:
+                            logger.warning("Empty search query in tool call")
+                            tool_outputs.append({
+                                "tool_call_id": tool_call.id,
+                                "output": "Error: Empty search query provided."
+                            })
+                            
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Failed to parse tool call arguments: {e}")
+                        tool_outputs.append({
+                            "tool_call_id": tool_call.id,
+                            "output": "Error: Failed to parse search parameters."
+                        })
+
+            # Continue conversation with search results
+            messages.append(llm_response.choices[0].message)
+            for output_item in tool_outputs:
+                messages.append({
+                    "role": "tool", 
+                    "tool_call_id": output_item["tool_call_id"], 
+                    "content": output_item["output"]
+                })
 
-        # Streaming generator
-        def stream_llm_response():
-            response = client.chat.completions.create(**llm_kwargs)
-            for chunk in response:
-                # Each chunk is an object, get the content delta
-                if hasattr(chunk.choices[0].delta, "content"):
-                    content = chunk.choices[0].delta.content
-                    if content:
-                        yield content
+            # Final response generation with search context
+            final_response = client.chat.completions.create(
+                model="unsloth/Qwen3-30B-A3B-GGUF",
+                temperature=temperature,  # Use same temperature for consistency
+                messages=messages,
+                max_tokens=2000
+            )
+            final_chatbot_response = final_response.choices[0].message.content
+        else:
+            final_chatbot_response = llm_response.choices[0].message.content
 
-        # Return as streaming response
-        return StreamingResponse(stream_llm_response(), media_type="text/plain")
+        # Enhanced response structure
+        response_data = {
+            "response": final_chatbot_response,
+            "sources": source_links,
+            "search_used": bool(tool_calls),
+            "temperature": temperature,  # Include temperature in response for debugging
+            "timestamp": datetime.now().isoformat()
+        }
+        
+        logger.info(f"Chat response generated successfully. Search used: {bool(tool_calls)}, Temperature: {temperature}")
+        return response_data
 
     except HTTPException:
         raise
@@ -276,7 +410,6 @@ async def chat_endpoint(request: Request):
         logger.error(f"Unexpected error in /chat endpoint: {e}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-
 # --- Health Check Endpoint ---
 @app.get("/")
 async def root():