Add Gemini model routing, tools, streaming, and doc-retrieval bridge

Author: owenisas

backend/agent.py

@@ -19,11 +19,12 @@
     openai_chat_stream,
     llm_chat,
     llm_chat_stream,
-    is_claude_model,
+    is_gemini_model,
     openai_upload_file,
     openai_create_vector_store,
     openai_wait_for_vector_store_files,
     openai_add_files_to_vector_store,
+    openai_search_vector_store,
 )
 from tutor_prompts import DEFAULT_USER_PROFILE
 from brave_images import BraveImageError, search_brave_images, verify_image_url
@@ -251,6 +252,32 @@
             },
         },
     },
+    {
+        "type": "function",
+        "function": {
+            "name": "search_uploaded_documents",
+            "description": (
+                "Search chunks from user-uploaded documents already indexed for this session. "
+                "Use this before asking the user to paste document content. "
+                "Best for questions about attached PDFs/docs."
+            ),
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The question or search query to run against uploaded documents.",
+                    },
+                    "max_results": {
+                        "type": "integer",
+                        "description": "Maximum number of chunks to return (1-10).",
+                        "default": 6,
+                    },
+                },
+                "required": ["query"],
+            },
+        },
+    },
     {
         "type": "function",
         "function": {
@@ -653,16 +680,22 @@ def _looks_like_file_access_failure(text):
     return any(token in cleaned for token in indicators)
 
 
-def _tools_for_request(web_search=False):
+def _tools_for_request(web_search=False, include_uploaded_documents_search=False):
     """
     Build tool list for this run.
 
-    Always returns the full tool set. Image-search function tools
+    Returns the tool set for this run. Image-search function tools
     (search_brave_images, search_wiki_images) must remain available even
     when web_search is enabled because OpenAI's built-in web_search does
     not return images.
     """
-    return TOOLS
+    if include_uploaded_documents_search:
+        return TOOLS
+    return [
+        tool
+        for tool in TOOLS
+        if str((tool.get("function") or {}).get("name") or "").strip() != "search_uploaded_documents"
+    ]
 
 
 def _append_reasoning_summary(block, reasoning):
@@ -804,6 +837,8 @@ def assess_context_or_answer_simple(
     system_prompt = get_pre_flight_prompt()
 
     _PRE_FLIGHT_TOOL_NAMES = {"update_learner_profile", "lookup_platform_help"}
+    if is_gemini_model(model) and enable_file_search and vector_store_id:
+        _PRE_FLIGHT_TOOL_NAMES.add("search_uploaded_documents")
     tools = [t for t in TOOLS if t["function"]["name"] in _PRE_FLIGHT_TOOL_NAMES]
 
     # Append current profile context
@@ -822,6 +857,16 @@ def assess_context_or_answer_simple(
                 ),
             }
         )
+        if is_gemini_model(model):
+            messages.append(
+                {
+                    "role": "system",
+                    "content": (
+                        "When document-specific evidence is needed, call search_uploaded_documents "
+                        "before asking the user for manual excerpts."
+                    ),
+                }
+            )
 
     if session_messages:
         # We append the conversation history, which already contains the user's latest message
@@ -1049,6 +1094,40 @@ def execute_tool(name, arguments, session=None):
             "current_profile": current_profile,
         }
 
+    elif name == "search_uploaded_documents":
+        query = str(arguments.get("query") or "").strip()
+        if not query:
+            return {"success": False, "error": "Missing required field: query"}
+
+        vector_store_id = str(arguments.get("vector_store_id") or "").strip()
+        if not vector_store_id and session is not None:
+            vector_store_id = str(getattr(session, "openai_vector_store_id", "") or "").strip()
+        if not vector_store_id:
+            return {"success": False, "error": "No uploaded-document index is available for this session."}
+
+        max_results = arguments.get("max_results", 6)
+        try:
+            max_results = int(max_results)
+        except (TypeError, ValueError):
+            max_results = 6
+        max_results = max(1, min(max_results, 10))
+
+        try:
+            payload = openai_search_vector_store(
+                vector_store_id=vector_store_id,
+                query=query,
+                max_results=max_results,
+            )
+            chunks = payload.get("results") if isinstance(payload, dict) else []
+            return {
+                "success": True,
+                "vector_store_id": vector_store_id,
+                "query": query,
+                "results": chunks if isinstance(chunks, list) else [],
+            }
+        except Exception as e:
+            return {"success": False, "error": f"Document search failed: {e}"}
+
     elif name == "create_main_block":
         title = str(arguments.get("title") or "").strip()
         summary = str(arguments.get("summary") or "").strip()
@@ -1589,6 +1668,7 @@ def run_agent_with_session(
     from models import Block
 
     model = model or os.getenv("OPENAI_MODEL", "gpt-5.2")
+    gemini_request = is_gemini_model(model)
 
     # Add user message to session, including attachment metadata for persistence.
     session.add_user_message(user_message, attachments=attachment_metadata)
@@ -1674,7 +1754,7 @@ def _mark_attachment(
     vector_store_id = getattr(session, "openai_vector_store_id", None)
     enable_file_search = False
 
-    if searchable_files and model.startswith("gpt"):
+    if searchable_files and (model.startswith("gpt") or gemini_request):
         upload_records = []
         for sf in searchable_files:
             filename = str(sf.get("filename") or "file")
@@ -1822,14 +1902,24 @@ def _mark_attachment(
 
     if vector_store_id:
         enable_file_search = True
-        _insert_system_instruction(
-            messages,
-            (
-                "File search context is available for this session. "
-                "When the user asks about uploaded documents, use file_search results first. "
-                "Do not ask the user to paste the document text unless retrieval fails."
-            ),
-        )
+        if gemini_request:
+            _insert_system_instruction(
+                messages,
+                (
+                    "Uploaded document retrieval is available via search_uploaded_documents. "
+                    "Use it before asking the user to paste document text. "
+                    "When citing results, include relevant excerpt details."
+                ),
+            )
+        else:
+            _insert_system_instruction(
+                messages,
+                (
+                    "File search context is available for this session. "
+                    "When the user asks about uploaded documents, use file_search results first. "
+                    "Do not ask the user to paste the document text unless retrieval fails."
+                ),
+            )
 
     # --- Pre-Flight Check ---
     # Only if this is a main block generation (depth 0)
@@ -1878,7 +1968,10 @@ def _mark_attachment(
             return response_block
 
     forced_tool_choice = None
-    active_tools = _tools_for_request(web_search)
+    active_tools = _tools_for_request(
+        web_search,
+        include_uploaded_documents_search=bool(gemini_request and vector_store_id),
+    )
     created_main_block = None
     should_stream = callable(stream_text_callback) or callable(stream_reasoning_callback)
 

backend/lambda_handler.py

@@ -143,8 +143,8 @@ def _emit_trace_event(trace_event: str, **fields) -> None:
 )
 
 # Allowed LLM models that users may select
-ALLOWED_MODELS = {"gpt-5.2", "claude-sonnet-4-6", "gemini-2.5-pro"}
-_ALLOWED_MODEL_PREFIXES = ()
+ALLOWED_MODELS = {"gpt-5.2", "claude-sonnet-4-6", "gemini-3.1-pro-preview"}
+_ALLOWED_MODEL_PREFIXES = ("gemini-",)
 
 
 def _is_allowed_model(model_name: str) -> bool: