Merge pull request #1130 from CREDO23/feat/vision-autocomplete

[Feat] Vision-based autocomplete with KB grounding

Author: MODSetter

surfsense_backend/alembic/versions/119_add_vision_llm_id_to_search_spaces.py

@@ -0,0 +1,39 @@
+"""119_add_vision_llm_id_to_search_spaces
+
+Revision ID: 119
+Revises: 118
+
+Adds vision_llm_id column to search_spaces for vision/screenshot analysis
+LLM role assignment. Defaults to 0 (Auto mode), same convention as
+agent_llm_id and document_summary_llm_id.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision: str = "119"
+down_revision: str | None = "118"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    existing_columns = [
+        col["name"] for col in sa.inspect(conn).get_columns("searchspaces")
+    ]
+
+    if "vision_llm_id" not in existing_columns:
+        op.add_column(
+            "searchspaces",
+            sa.Column("vision_llm_id", sa.Integer(), nullable=True, server_default="0"),
+        )
+
+
+def downgrade() -> None:
+    op.drop_column("searchspaces", "vision_llm_id")

surfsense_backend/app/db.py

@@ -1351,6 +1351,9 @@ class SearchSpace(BaseModel, TimestampMixin):
     image_generation_config_id = Column(
         Integer, nullable=True, default=0
     )  # For image generation, defaults to Auto mode
+    vision_llm_id = Column(
+        Integer, nullable=True, default=0
+    )  # For vision/screenshot analysis, defaults to Auto mode
 
     user_id = Column(
         UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False

surfsense_backend/app/routes/__init__.py

@@ -3,6 +3,7 @@
 from .airtable_add_connector_route import (
     router as airtable_add_connector_router,
 )
+from .autocomplete_routes import router as autocomplete_router
 from .chat_comments_routes import router as chat_comments_router
 from .circleback_webhook_route import router as circleback_webhook_router
 from .clickup_add_connector_route import router as clickup_add_connector_router
@@ -95,3 +96,4 @@
 router.include_router(stripe_router)  # Stripe checkout for additional page packs
 router.include_router(youtube_router)  # YouTube playlist resolution
 router.include_router(prompts_router)
+router.include_router(autocomplete_router)  # Lightweight autocomplete with KB context

surfsense_backend/app/routes/autocomplete_routes.py

@@ -0,0 +1,42 @@
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db import User, get_async_session
+from app.services.new_streaming_service import VercelStreamingService
+from app.services.vision_autocomplete_service import stream_vision_autocomplete
+from app.users import current_active_user
+from app.utils.rbac import check_search_space_access
+
+router = APIRouter(prefix="/autocomplete", tags=["autocomplete"])
+
+MAX_SCREENSHOT_SIZE = 20 * 1024 * 1024  # 20 MB base64 ceiling
+
+
+class VisionAutocompleteRequest(BaseModel):
+    screenshot: str = Field(..., max_length=MAX_SCREENSHOT_SIZE)
+    search_space_id: int
+    app_name: str = ""
+    window_title: str = ""
+
+
+@router.post("/vision/stream")
+async def vision_autocomplete_stream(
+    body: VisionAutocompleteRequest,
+    user: User = Depends(current_active_user),
+    session: AsyncSession = Depends(get_async_session),
+):
+    await check_search_space_access(session, user, body.search_space_id)
+
+    return StreamingResponse(
+        stream_vision_autocomplete(
+            body.screenshot, body.search_space_id, session,
+            app_name=body.app_name, window_title=body.window_title,
+        ),
+        media_type="text/event-stream",
+        headers={
+            **VercelStreamingService.get_response_headers(),
+            "X-Accel-Buffering": "no",
+        },
+    )

surfsense_backend/app/routes/search_spaces_routes.py

@@ -522,14 +522,17 @@ async def get_llm_preferences(
         image_generation_config = await _get_image_gen_config_by_id(
             session, search_space.image_generation_config_id
         )
+        vision_llm = await _get_llm_config_by_id(session, search_space.vision_llm_id)
 
         return LLMPreferencesRead(
             agent_llm_id=search_space.agent_llm_id,
             document_summary_llm_id=search_space.document_summary_llm_id,
             image_generation_config_id=search_space.image_generation_config_id,
+            vision_llm_id=search_space.vision_llm_id,
             agent_llm=agent_llm,
             document_summary_llm=document_summary_llm,
             image_generation_config=image_generation_config,
+            vision_llm=vision_llm,
         )
 
     except HTTPException:
@@ -589,14 +592,17 @@ async def update_llm_preferences(
         image_generation_config = await _get_image_gen_config_by_id(
             session, search_space.image_generation_config_id
         )
+        vision_llm = await _get_llm_config_by_id(session, search_space.vision_llm_id)
 
         return LLMPreferencesRead(
             agent_llm_id=search_space.agent_llm_id,
             document_summary_llm_id=search_space.document_summary_llm_id,
             image_generation_config_id=search_space.image_generation_config_id,
+            vision_llm_id=search_space.vision_llm_id,
             agent_llm=agent_llm,
             document_summary_llm=document_summary_llm,
             image_generation_config=image_generation_config,
+            vision_llm=vision_llm,
         )
 
     except HTTPException:

surfsense_backend/app/schemas/new_llm_config.py

@@ -182,6 +182,9 @@ class LLMPreferencesRead(BaseModel):
     image_generation_config_id: int | None = Field(
         None, description="ID of the image generation config to use"
     )
+    vision_llm_id: int | None = Field(
+        None, description="ID of the LLM config to use for vision/screenshot analysis"
+    )
     agent_llm: dict[str, Any] | None = Field(
         None, description="Full config for agent LLM"
     )
@@ -191,6 +194,9 @@ class LLMPreferencesRead(BaseModel):
     image_generation_config: dict[str, Any] | None = Field(
         None, description="Full config for image generation"
     )
+    vision_llm: dict[str, Any] | None = Field(
+        None, description="Full config for vision LLM"
+    )
 
     model_config = ConfigDict(from_attributes=True)
 
@@ -207,3 +213,6 @@ class LLMPreferencesUpdate(BaseModel):
     image_generation_config_id: int | None = Field(
         None, description="ID of the image generation config to use"
     )
+    vision_llm_id: int | None = Field(
+        None, description="ID of the LLM config to use for vision/screenshot analysis"
+    )

surfsense_backend/app/services/llm_service.py

@@ -32,6 +32,7 @@
 class LLMRole:
     AGENT = "agent"  # For agent/chat operations
     DOCUMENT_SUMMARY = "document_summary"  # For document summarization
+    VISION = "vision"  # For vision/screenshot analysis
 
 
 def get_global_llm_config(llm_config_id: int) -> dict | None:
@@ -187,7 +188,7 @@ async def get_search_space_llm_instance(
     Args:
         session: Database session
         search_space_id: Search Space ID
-        role: LLM role ('agent' or 'document_summary')
+        role: LLM role ('agent', 'document_summary', or 'vision')
 
     Returns:
         ChatLiteLLM or ChatLiteLLMRouter instance, or None if not found
@@ -209,6 +210,8 @@ async def get_search_space_llm_instance(
             llm_config_id = search_space.agent_llm_id
         elif role == LLMRole.DOCUMENT_SUMMARY:
             llm_config_id = search_space.document_summary_llm_id
+        elif role == LLMRole.VISION:
+            llm_config_id = search_space.vision_llm_id
         else:
             logger.error(f"Invalid LLM role: {role}")
             return None
@@ -405,6 +408,13 @@ async def get_document_summary_llm(
     )
 
 
+async def get_vision_llm(
+    session: AsyncSession, search_space_id: int
+) -> ChatLiteLLM | ChatLiteLLMRouter | None:
+    """Get the search space's vision LLM instance for screenshot analysis."""
+    return await get_search_space_llm_instance(session, search_space_id, LLMRole.VISION)
+
+
 # Backward-compatible alias (LLM preferences are now per-search-space, not per-user)
 async def get_user_long_context_llm(
     session: AsyncSession,