diff --git a/surfsense_backend/app/app.py b/surfsense_backend/app/app.py index 8d525d659..3fc0b14f7 100644 --- a/surfsense_backend/app/app.py +++ b/surfsense_backend/app/app.py @@ -36,7 +36,7 @@ app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*") # Add CORS middleware app.add_middleware( CORSMiddleware, - allow_origins=["*"], # Allows all origins + allow_origins=["https://ai.kapteinis.lv"], # Only allow our domain allow_credentials=True, allow_methods=["*"], # Allows all methods allow_headers=["*"], # Allows all headers diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index c23c0d13e..ce386e554 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -1,9 +1,10 @@ from collections.abc import AsyncGenerator from datetime import UTC, datetime from enum import Enum +import uuid from fastapi import Depends -from fastapi_users.db import SQLAlchemyBaseUserTableUUID, SQLAlchemyUserDatabase +from fastapi_users_db_sqlalchemy import SQLAlchemyBaseUserTable, SQLAlchemyUserDatabase from pgvector.sqlalchemy import Vector from sqlalchemy import ( ARRAY, @@ -28,9 +29,6 @@ from app.config import config from app.retriver.chunks_hybrid_search import ChucksHybridSearchRetriever from app.retriver.documents_hybrid_search import DocumentHybridSearchRetriever -if config.AUTH_TYPE == "GOOGLE": - from fastapi_users.db import SQLAlchemyBaseOAuthAccountTableUUID - DATABASE_URL = config.DATABASE_URL @@ -55,11 +53,11 @@ class DocumentType(str, Enum): class SearchSourceConnectorType(str, Enum): - SERPER_API = "SERPER_API" # NOT IMPLEMENTED YET : DON'T REMEMBER WHY : MOST PROBABLY BECAUSE WE NEED TO CRAWL THE RESULTS RETURNED BY IT + SERPER_API = "SERPER_API" TAVILY_API = "TAVILY_API" SEARXNG_API = "SEARXNG_API" LINKUP_API = "LINKUP_API" - BAIDU_SEARCH_API = "BAIDU_SEARCH_API" # Baidu AI Search API for Chinese web search + BAIDU_SEARCH_API = "BAIDU_SEARCH_API" SLACK_CONNECTOR = "SLACK_CONNECTOR" NOTION_CONNECTOR = "NOTION_CONNECTOR" GITHUB_CONNECTOR = "GITHUB_CONNECTOR" @@ -80,10 +78,6 @@ class ChatType(str, Enum): class LiteLLMProvider(str, Enum): - """ - Enum for LLM providers supported by LiteLLM. - """ - OPENAI = "OPENAI" ANTHROPIC = "ANTHROPIC" GOOGLE = "GOOGLE" @@ -136,7 +130,7 @@ class Base(DeclarativeBase): class TimestampMixin: @declared_attr - def created_at(cls): # noqa: N805 + def created_at(cls): return Column( TIMESTAMP(timezone=True), nullable=False, @@ -208,7 +202,7 @@ class Podcast(BaseModel, TimestampMixin): file_location = Column(String(500), nullable=False, default="") chat_id = Column( Integer, ForeignKey("chats.id", ondelete="CASCADE"), nullable=True - ) # If generated from a chat, this will be the chat id, else null ( can be from a document or a chat ) + ) chat_state_version = Column(BigInteger, nullable=True) search_space_id = Column( @@ -288,7 +282,6 @@ class SearchSourceConnector(BaseModel, TimestampMixin): last_indexed_at = Column(TIMESTAMP(timezone=True), nullable=True) config = Column(JSON, nullable=False) - # Periodic indexing fields periodic_indexing_enabled = Column(Boolean, nullable=False, default=False) indexing_frequency_minutes = Column(Integer, nullable=True) next_scheduled_at = Column(TIMESTAMP(timezone=True), nullable=True) @@ -309,19 +302,14 @@ class LLMConfig(BaseModel, TimestampMixin): __tablename__ = "llm_configs" name = Column(String(100), nullable=False, index=True) - # Provider from the enum provider = Column(SQLAlchemyEnum(LiteLLMProvider), nullable=False) - # Custom provider name when provider is CUSTOM custom_provider = Column(String(100), nullable=True) - # Just the model name without provider prefix model_name = Column(String(100), nullable=False) - # API Key should be encrypted before storing api_key = Column(String, nullable=False) api_base = Column(String(500), nullable=True) language = Column(String(50), nullable=True, default="English") - # For any other parameters that litellm supports litellm_params = Column(JSON, nullable=True, default={}) search_space_id = Column( @@ -347,27 +335,13 @@ class UserSearchSpacePreference(BaseModel, TimestampMixin): Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False ) - # User-specific LLM preferences for this search space - # Note: These can be negative IDs for global configs (from YAML) or positive IDs for custom configs (from DB) - # Foreign keys removed to support global configs with negative IDs long_context_llm_id = Column(Integer, nullable=True) fast_llm_id = Column(Integer, nullable=True) strategic_llm_id = Column(Integer, nullable=True) - # Future RBAC fields can be added here - # role = Column(String(50), nullable=True) # e.g., 'owner', 'editor', 'viewer' - # permissions = Column(JSON, nullable=True) - user = relationship("User", back_populates="search_space_preferences") search_space = relationship("SearchSpace", back_populates="user_preferences") - # Note: Relationships removed because foreign keys no longer exist - # Global configs (negative IDs) don't exist in llm_configs table - # Application code manually fetches configs when needed - # long_context_llm = relationship("LLMConfig", foreign_keys=[long_context_llm_id], post_update=True) - # fast_llm = relationship("LLMConfig", foreign_keys=[fast_llm_id], post_update=True) - # strategic_llm = relationship("LLMConfig", foreign_keys=[strategic_llm_id], post_update=True) - class Log(BaseModel, TimestampMixin): __tablename__ = "logs" @@ -375,10 +349,8 @@ class Log(BaseModel, TimestampMixin): level = Column(SQLAlchemyEnum(LogLevel), nullable=False, index=True) status = Column(SQLAlchemyEnum(LogStatus), nullable=False, index=True) message = Column(Text, nullable=False) - source = Column( - String(200), nullable=True, index=True - ) # Service/component that generated the log - log_metadata = Column(JSON, nullable=True, default={}) # Additional context data + source = Column(String(200), nullable=True, index=True) + log_metadata = Column(JSON, nullable=True, default={}) search_space_id = Column( Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False @@ -388,36 +360,29 @@ class Log(BaseModel, TimestampMixin): if config.AUTH_TYPE == "GOOGLE": - class OAuthAccount(SQLAlchemyBaseOAuthAccountTableUUID, Base): - pass + class OAuthAccount(Base): + __tablename__ = "oauth_account" + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) - class User(SQLAlchemyBaseUserTableUUID, Base): - oauth_accounts: Mapped[list[OAuthAccount]] = relationship( - "OAuthAccount", lazy="joined" - ) + class User(SQLAlchemyBaseUserTable, Base): + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + oauth_accounts = relationship("OAuthAccount", lazy="joined") search_spaces = relationship("SearchSpace", back_populates="user") search_space_preferences = relationship( - "UserSearchSpacePreference", - back_populates="user", - cascade="all, delete-orphan", + "UserSearchSpacePreference", back_populates="user", cascade="all, delete-orphan" ) - - # Page usage tracking for ETL services - pages_limit = Column(Integer, nullable=False, default=500, server_default="500") + pages_limit = Column(Integer, nullable=False, default=1000, server_default="1000") pages_used = Column(Integer, nullable=False, default=0, server_default="0") else: - class User(SQLAlchemyBaseUserTableUUID, Base): + class User(SQLAlchemyBaseUserTable, Base): + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) search_spaces = relationship("SearchSpace", back_populates="user") search_space_preferences = relationship( - "UserSearchSpacePreference", - back_populates="user", - cascade="all, delete-orphan", + "UserSearchSpacePreference", back_populates="user", cascade="all, delete-orphan" ) - - # Page usage tracking for ETL services - pages_limit = Column(Integer, nullable=False, default=500, server_default="500") + pages_limit = Column(Integer, nullable=False, default=1000, server_default="1000") pages_used = Column(Integer, nullable=False, default=0, server_default="0") @@ -427,8 +392,6 @@ async_session_maker = async_sessionmaker(engine, expire_on_commit=False) async def setup_indexes(): async with engine.begin() as conn: - # Create indexes - # Document Summary Indexes await conn.execute( text( "CREATE INDEX IF NOT EXISTS document_vector_index ON documents USING hnsw (embedding public.vector_cosine_ops)" @@ -439,7 +402,6 @@ async def setup_indexes(): "CREATE INDEX IF NOT EXISTS document_search_index ON documents USING gin (to_tsvector('english', content))" ) ) - # Document Chuck Indexes await conn.execute( text( "CREATE INDEX IF NOT EXISTS chucks_vector_index ON chunks USING hnsw (embedding public.vector_cosine_ops)" @@ -485,3 +447,4 @@ async def get_documents_hybrid_search_retriever( session: AsyncSession = Depends(get_async_session), ): return DocumentHybridSearchRetriever(session) + diff --git a/surfsense_backend/app/routes/documents_routes.py b/surfsense_backend/app/routes/documents_routes.py index 344a2503d..4065deeb4 100644 --- a/surfsense_backend/app/routes/documents_routes.py +++ b/surfsense_backend/app/routes/documents_routes.py @@ -108,17 +108,21 @@ async def create_documents_file_upload( for file in files: try: - # Save file to a temporary location to avoid stream issues + # Save file to persistent uploads directory import os - import tempfile + import uuid + from pathlib import Path - # Create temp file - with tempfile.NamedTemporaryFile( - delete=False, suffix=os.path.splitext(file.filename)[1] - ) as temp_file: - temp_path = temp_file.name + # Create uploads directory if it doesn't exist + uploads_dir = Path("/opt/SurfSense/surfsense_backend/uploads") + uploads_dir.mkdir(parents=True, exist_ok=True) - # Write uploaded file to temp file + # Create unique filename + file_ext = os.path.splitext(file.filename)[1] + unique_filename = f"{uuid.uuid4()}{file_ext}" + temp_path = str(uploads_dir / unique_filename) + + # Write uploaded file to persistent location content = await file.read() with open(temp_path, "wb") as f: f.write(content) diff --git a/surfsense_backend/app/services/streaming_service.py b/surfsense_backend/app/services/streaming_service.py index 98c0d3ac5..389380893 100644 --- a/surfsense_backend/app/services/streaming_service.py +++ b/surfsense_backend/app/services/streaming_service.py @@ -189,3 +189,19 @@ class StreamingService: }, } return f"d:{json.dumps(completion_data)}\n" + + def format_grammar_check_delta(self, grammar_result: dict) -> str: + """ + Format grammar check results as a delta annotation + + Args: + grammar_result: Grammar check result dictionary + + Returns: + str: The formatted annotation delta string + """ + annotation = { + "type": "GRAMMAR_CHECK", + "data": grammar_result, + } + return f"8:[{json.dumps(annotation)}]\n" diff --git a/surfsense_backend/app/tasks/stream_connector_search_results.py b/surfsense_backend/app/tasks/stream_connector_search_results.py index a944b1cfd..8e8fedf60 100644 --- a/surfsense_backend/app/tasks/stream_connector_search_results.py +++ b/surfsense_backend/app/tasks/stream_connector_search_results.py @@ -1,4 +1,8 @@ from collections.abc import AsyncGenerator +import asyncio +import json +import logging +import os from typing import Any from uuid import UUID @@ -8,6 +12,9 @@ from app.agents.researcher.configuration import SearchMode from app.agents.researcher.graph import graph as researcher_graph from app.agents.researcher.state import State from app.services.streaming_service import StreamingService +from app.services.grammar_check import auto_grammar_check + +logger = logging.getLogger(__name__) async def stream_connector_search_results( @@ -71,6 +78,9 @@ async def stream_connector_search_results( # Run the graph directly print("\nRunning the complete researcher workflow...") + # Collect response text for grammar checking + response_chunks = [] + # Use streaming with config parameter async for chunk in researcher_graph.astream( initial_state, @@ -78,6 +88,45 @@ async def stream_connector_search_results( stream_mode="custom", ): if isinstance(chunk, dict) and "yield_value" in chunk: - yield chunk["yield_value"] + # Collect text chunks for grammar checking + yield_value = chunk["yield_value"] + yield yield_value + + # Try to extract text from the chunk + try: + # Parse the chunk to extract text + # Format is like "0:"text"\n" for text chunks + if yield_value.startswith("0:"): + text_json = yield_value[2:].strip() + if text_json: + text = json.loads(text_json) + response_chunks.append(text) + except Exception: + # Ignore parsing errors, not all chunks contain text + pass + + # Run grammar check asynchronously with timeout + ollama_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") + + try: + # Combine response chunks + full_response = "".join(response_chunks) + + if full_response.strip(): + # Run grammar check with timeout + grammar_result = await asyncio.wait_for( + auto_grammar_check(user_query, full_response, ollama_url), + timeout=8.0 + ) + + # If grammar check returned a result, stream it + if grammar_result: + logger.info(f"Grammar check completed for language: {grammar_result.get('language_code', 'unknown')}") + yield streaming_service.format_grammar_check_delta(grammar_result) + + except asyncio.TimeoutError: + logger.warning("Grammar check timed out") + except Exception as e: + logger.warning(f"Grammar check failed: {e}") yield streaming_service.format_completion() diff --git a/surfsense_backend/app/users.py b/surfsense_backend/app/users.py index d51b30bd7..e46799186 100644 --- a/surfsense_backend/app/users.py +++ b/surfsense_backend/app/users.py @@ -8,7 +8,7 @@ from fastapi_users.authentication import ( BearerTransport, JWTStrategy, ) -from fastapi_users.db import SQLAlchemyUserDatabase +from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase from pydantic import BaseModel from app.config import config