apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-07-22 23:31:12 +02:00
Code Issues Projects Releases Packages Wiki Activity

Update streaming service and database connection handling

Browse source 
- Refactor database connection management in db.py
- Enhance streaming service with additional features
- Update document routes for improved performance
- Improve stream connector search results processing
- Update app.py and users.py configurations

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Ojārs Kapteinis 2025-11-17 20:34:56 +02:00
parent fdef50e78d
commit 1378828b03
6 changed files with 101 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

2
surfsense_backend/app/app.py
View file

@ -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

79
surfsense_backend/app/db.py
View file

@ -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)

20
surfsense_backend/app/routes/documents_routes.py
View file

@ -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)

16
surfsense_backend/app/services/streaming_service.py
View file

@ -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"

51
surfsense_backend/app/tasks/stream_connector_search_results.py
View file

@ -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()

2
surfsense_backend/app/users.py
View file

@ -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