feat: migrated old chat to new chat

surfsense_backend/alembic/versions/49_migrate_old_chats_to_new_chat.py

@@ -0,0 +1,216 @@
+"""Migrate old chats to new_chat_threads and remove old tables
+
+Revision ID: 49
+Revises: 48
+Create Date: 2025-12-21
+
+This migration:
+1. Migrates data from old 'chats' table to 'new_chat_threads' and 'new_chat_messages'
+2. Drops the 'podcasts' table (podcast data is not migrated as per user request)
+3. Drops the 'chats' table
+4. Removes the 'chattype' enum
+"""
+
+import json
+from collections.abc import Sequence
+from datetime import datetime, timezone
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "49"
+down_revision: str | None = "48"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def extract_text_content(content: str | dict | list) -> str:
+    """Extract plain text content from various message formats."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, dict):
+        # Handle dict with 'text' key
+        if "text" in content:
+            return content["text"]
+        return str(content)
+    if isinstance(content, list):
+        # Handle list of parts (e.g., [{"type": "text", "text": "..."}])
+        texts = []
+        for part in content:
+            if isinstance(part, dict) and part.get("type") == "text":
+                texts.append(part.get("text", ""))
+            elif isinstance(part, str):
+                texts.append(part)
+        return "\n".join(texts) if texts else ""
+    return ""
+
+
+def parse_timestamp(ts, fallback):
+    """Parse ISO timestamp string to datetime object."""
+    if ts is None:
+        return fallback
+    if isinstance(ts, datetime):
+        return ts
+    if isinstance(ts, str):
+        try:
+            # Handle ISO format like '2025-11-26T22:43:34.399Z'
+            ts = ts.replace("Z", "+00:00")
+            return datetime.fromisoformat(ts)
+        except (ValueError, TypeError):
+            return fallback
+    return fallback
+
+
+def upgrade() -> None:
+    """Migrate old chats to new_chat_threads and remove old tables."""
+    connection = op.get_bind()
+
+    # Get all old chats
+    old_chats = connection.execute(
+        sa.text("""
+            SELECT id, title, messages, search_space_id, created_at
+            FROM chats
+            ORDER BY created_at ASC
+        """)
+    ).fetchall()
+
+    print(f"[Migration 49] Found {len(old_chats)} old chats to migrate")
+
+    migrated_count = 0
+    for chat_id, title, messages_json, search_space_id, created_at in old_chats:
+        try:
+            # Parse messages JSON
+            if isinstance(messages_json, str):
+                messages = json.loads(messages_json)
+            else:
+                messages = messages_json or []
+
+            # Skip empty chats
+            if not messages:
+                print(f"[Migration 49] Skipping empty chat {chat_id}")
+                continue
+
+            # Create new thread
+            result = connection.execute(
+                sa.text("""
+                    INSERT INTO new_chat_threads 
+                    (title, archived, search_space_id, created_at, updated_at)
+                    VALUES (:title, FALSE, :search_space_id, :created_at, :created_at)
+                    RETURNING id
+                """),
+                {
+                    "title": title or "Migrated Chat",
+                    "search_space_id": search_space_id,
+                    "created_at": created_at,
+                },
+            )
+            new_thread_id = result.fetchone()[0]
+
+            # Migrate messages - only user and assistant roles, skip SOURCES/TERMINAL_INFO
+            message_count = 0
+            for msg in messages:
+                role_lower = msg.get("role", "").lower()
+
+                # Only migrate user and assistant messages
+                if role_lower not in ("user", "assistant"):
+                    continue
+                
+                # Convert to uppercase for database enum
+                role = role_lower.upper()
+
+                # Extract content - handle various formats
+                content_raw = msg.get("content", "")
+                content_text = extract_text_content(content_raw)
+
+                # Skip empty messages
+                if not content_text.strip():
+                    continue
+
+                # Parse message timestamp
+                msg_created_at = parse_timestamp(msg.get("createdAt"), created_at)
+
+                # Store content as JSONB array format for assistant-ui compatibility
+                content_list = [{"type": "text", "text": content_text}]
+
+                # Use direct SQL with string interpolation for the enum since CAST doesn't work
+                # The enum value comes from trusted source (our own code), not user input
+                connection.execute(
+                    sa.text(f"""
+                        INSERT INTO new_chat_messages 
+                        (thread_id, role, content, created_at)
+                        VALUES (:thread_id, '{role}', CAST(:content AS jsonb), :created_at)
+                    """),
+                    {
+                        "thread_id": new_thread_id,
+                        "content": json.dumps(content_list),
+                        "created_at": msg_created_at,
+                    },
+                )
+                message_count += 1
+
+            print(
+                f"[Migration 49] Migrated chat {chat_id} -> thread {new_thread_id} ({message_count} messages)"
+            )
+            migrated_count += 1
+
+        except Exception as e:
+            print(f"[Migration 49] Error migrating chat {chat_id}: {e}")
+            # Re-raise to abort migration - we don't want partial data
+            raise
+
+    print(f"[Migration 49] Successfully migrated {migrated_count} chats")
+
+    # Drop podcasts table (FK references chats, so drop first)
+    print("[Migration 49] Dropping podcasts table...")
+    op.drop_table("podcasts")
+
+    # Drop chats table
+    print("[Migration 49] Dropping chats table...")
+    op.drop_table("chats")
+
+    # Drop chattype enum
+    print("[Migration 49] Dropping chattype enum...")
+    op.execute(sa.text("DROP TYPE IF EXISTS chattype"))
+
+    print("[Migration 49] Migration complete!")
+
+
+def downgrade() -> None:
+    """Recreate old tables (data cannot be restored)."""
+    # Recreate chattype enum
+    op.execute(
+        sa.text("""
+            CREATE TYPE chattype AS ENUM ('QNA')
+        """)
+    )
+
+    # Recreate chats table
+    op.create_table(
+        "chats",
+        sa.Column("id", sa.Integer(), primary_key=True, index=True),
+        sa.Column("type", sa.Enum("QNA", name="chattype"), nullable=False),
+        sa.Column("title", sa.String(), nullable=False, index=True),
+        sa.Column("initial_connectors", sa.ARRAY(sa.String()), nullable=True),
+        sa.Column("messages", sa.JSON(), nullable=False),
+        sa.Column("state_version", sa.BigInteger(), nullable=False, default=1),
+        sa.Column("search_space_id", sa.Integer(), sa.ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("created_at", sa.TIMESTAMP(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    # Recreate podcasts table
+    op.create_table(
+        "podcasts",
+        sa.Column("id", sa.Integer(), primary_key=True, index=True),
+        sa.Column("title", sa.String(), nullable=False, index=True),
+        sa.Column("podcast_transcript", sa.JSON(), nullable=False, server_default="{}"),
+        sa.Column("file_location", sa.String(500), nullable=False, server_default=""),
+        sa.Column("chat_id", sa.Integer(), sa.ForeignKey("chats.id", ondelete="CASCADE"), nullable=True),
+        sa.Column("chat_state_version", sa.BigInteger(), nullable=True),
+        sa.Column("search_space_id", sa.Integer(), sa.ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("created_at", sa.TIMESTAMP(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    print("[Migration 49 Downgrade] Tables recreated (data not restored)")
+

surfsense_backend/app/agents/researcher/configuration.py

@@ -1,30 +0,0 @@
-"""Define the configurable parameters for the agent."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, fields
-
-from langchain_core.runnables import RunnableConfig
-
-
-@dataclass(kw_only=True)
-class Configuration:
-    """The configuration for the agent."""
-
-    # Input parameters provided at invocation
-    user_query: str
-    connectors_to_search: list[str]
-    user_id: str
-    search_space_id: int
-    document_ids_to_add_in_context: list[int]
-    language: str | None = None
-    top_k: int = 10
-
-    @classmethod
-    def from_runnable_config(
-        cls, config: RunnableConfig | None = None
-    ) -> Configuration:
-        """Create a Configuration instance from a RunnableConfig object."""
-        configurable = (config.get("configurable") or {}) if config else {}
-        _fields = {f.name for f in fields(cls) if f.init}
-        return cls(**{k: v for k, v in configurable.items() if k in _fields})

surfsense_backend/app/agents/researcher/graph.py

@@ -1,47 +0,0 @@
-from langgraph.graph import StateGraph
-
-from .configuration import Configuration
-from .nodes import (
-    generate_further_questions,
-    handle_qna_workflow,
-    reformulate_user_query,
-)
-from .state import State
-
-
-def build_graph():
-    """
-    Build and return the LangGraph workflow.
-
-    This function constructs the researcher agent graph for Q&A workflow.
-    The workflow follows a simple path:
-    1. Reformulate user query based on chat history
-    2. Handle QNA workflow (fetch documents and generate answer)
-    3. Generate follow-up questions
-
-    Returns:
-        A compiled LangGraph workflow
-    """
-    # Define a new graph with state class
-    workflow = StateGraph(State, config_schema=Configuration)
-
-    # Add nodes to the graph
-    workflow.add_node("reformulate_user_query", reformulate_user_query)
-    workflow.add_node("handle_qna_workflow", handle_qna_workflow)
-    workflow.add_node("generate_further_questions", generate_further_questions)
-
-    # Define the edges - simple linear flow for QNA
-    workflow.add_edge("__start__", "reformulate_user_query")
-    workflow.add_edge("reformulate_user_query", "handle_qna_workflow")
-    workflow.add_edge("handle_qna_workflow", "generate_further_questions")
-    workflow.add_edge("generate_further_questions", "__end__")
-
-    # Compile the workflow into an executable graph
-    graph = workflow.compile()
-    graph.name = "Surfsense Researcher"  # This defines the custom name in LangSmith
-
-    return graph
-
-
-# Compile the graph once when the module is loaded
-graph = build_graph()

surfsense_backend/app/agents/researcher/prompts.py

@@ -1,140 +0,0 @@
-import datetime
-
-
-def _build_language_instruction(language: str | None = None):
-    """Build language instruction for prompts."""
-    if language:
-        return f"\n\nIMPORTANT: Please respond in {language} language. All your responses, explanations, and analysis should be written in {language}."
-    return ""
-
-
-def get_further_questions_system_prompt():
-    return f"""
-Today's date: {datetime.datetime.now().strftime("%Y-%m-%d")}
-<further_questions_system>
-You are an expert research assistant specializing in generating contextually relevant follow-up questions. Your task is to analyze the chat history and available documents to suggest further questions that would naturally extend the conversation and provide additional value to the user.
-
-<input>
-- chat_history: Provided in XML format within <chat_history> tags, containing <user> and <assistant> message pairs that show the chronological conversation flow. This provides context about what has already been discussed.
-- available_documents: Provided in XML format within <documents> tags, containing individual <document> elements with <document_metadata> and <document_content> sections. Each document contains multiple `<chunk id='...'>...</chunk>` blocks inside <document_content>. This helps understand what information is accessible for answering potential follow-up questions.
-</input>
-
-<output_format>
-A JSON object with the following structure:
-{{
-  "further_questions": [
-    {{
-      "id": 0,
-      "question": "further qn 1"
-    }},
-    {{
-      "id": 1,
-      "question": "further qn 2"
-    }}
-  ]
-}}
-</output_format>
-
-<instructions>
-1.  **Analyze Chat History:** Review the entire conversation flow to understand:
-    *   The main topics and themes discussed
-    *   The user's interests and areas of focus
-    *   Questions that have been asked and answered
-    *   Any gaps or areas that could be explored further
-    *   The depth level of the current discussion
-
-2.  **Evaluate Available Documents:** Consider the documents in context to identify:
-    *   Additional information that hasn't been explored yet
-    *   Related topics that could be of interest
-    *   Specific details or data points that could warrant deeper investigation
-    *   Cross-references or connections between different documents
-
-3.  **Generate Relevant Follow-up Questions:** Create 3-5 further questions that:
-    *   Are directly related to the ongoing conversation but explore new angles
-    *   Can be reasonably answered using the available documents or knowledge base
-    *   Progress the conversation forward rather than repeating previous topics
-    *   Match the user's apparent level of interest and expertise
-    *   Are specific and actionable rather than overly broad
-    *   Consider practical applications, comparisons, deeper analysis, or related concepts
-
-4.  **Ensure Question Quality:** Each question should:
-    *   Be clear and well-formulated
-    *   Provide genuine value to the user
-    *   Be distinct from other suggested questions
-    *   Be answerable within the current context
-    *   Encourage meaningful exploration of the topic
-
-5.  **Prioritize and Order:** Arrange questions by relevance and natural progression:
-    *   Most directly related to the current discussion first
-    *   Questions that build upon previous answers
-    *   Questions that explore practical applications or implications
-    *   Questions that introduce related but new concepts
-
-6.  **Adhere Strictly to Output Format:** Ensure the final output is a valid JSON object with:
-    *   Correct field names (`further_questions`, `id`, `question`)
-    *   Sequential numbering starting from 0
-    *   Proper data types and JSON formatting
-</instructions>
-
-<examples>
-Input:
-<chat_history>
-<user>What are the best machine learning algorithms for text classification?</user>
-<assistant>For text classification, several algorithms work well depending on your specific needs:
-
-**Traditional Methods:**
-- **Support Vector Machines (SVM)** - Excellent for high-dimensional text data
-- **Naive Bayes** - Simple, fast, and works well with small datasets
-- **Logistic Regression** - Good baseline with interpretable results
-
-**Modern Deep Learning:**
-- **Neural Networks** - More complex but can capture intricate patterns
-- **Transformer models** - State-of-the-art for most text classification tasks
-
-The choice depends on your dataset size, computational resources, and accuracy requirements.</assistant>
-</chat_history>
-
-<documents>
-<document>
-<metadata>
-<source_id>101</source_id>
-<source_type>FILE</source_type>
-</metadata>
-<content>
-# Machine Learning for Text Classification: A Comprehensive Guide
-
-## Performance Comparison
-Recent studies show that transformer-based models achieve 95%+ accuracy on most text classification benchmarks, while traditional methods like SVM typically achieve 85-90% accuracy.
-
-## Dataset Considerations
-- Small datasets (< 1000 samples): Naive Bayes, SVM
-- Large datasets (> 10,000 samples): Neural networks, transformers
-- Imbalanced datasets: Require special handling with techniques like SMOTE
-</content>
-</document>
-</documents>
-
-Output:
-{{
-  "further_questions": [
-    {{
-      "id": 0,
-      "question": "What are the key differences in performance between traditional algorithms like SVM and modern deep learning approaches for text classification?"
-    }},
-    {{
-      "id": 1,
-      "question": "How do you handle imbalanced datasets when training text classification models?"
-    }},
-    {{
-      "id": 2,
-      "question": "What preprocessing techniques are most effective for improving text classification accuracy?"
-    }},
-    {{
-      "id": 3,
-      "question": "Are there specific domains or use cases where certain classification algorithms perform better than others?"
-    }}
-  ]
-}}
-</examples>
-</further_questions_system>
-"""

surfsense_backend/app/agents/researcher/qna_agent/__init__.py

@@ -1,5 +0,0 @@
-"""QnA Agent."""
-
-from .graph import graph
-
-__all__ = ["graph"]

surfsense_backend/app/agents/researcher/qna_agent/configuration.py

@@ -1,31 +0,0 @@
-"""Define the configurable parameters for the agent."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, fields
-from typing import Any
-
-from langchain_core.runnables import RunnableConfig
-
-
-@dataclass(kw_only=True)
-class Configuration:
-    """The configuration for the Q&A agent."""
-
-    # Configuration parameters for the Q&A agent
-    user_query: str  # The user's question to answer
-    reformulated_query: str  # The reformulated query
-    relevant_documents: list[
-        Any
-    ]  # Documents provided directly to the agent for answering
-    search_space_id: int  # Search space identifier
-    language: str | None = None  # Language for responses
-
-    @classmethod
-    def from_runnable_config(
-        cls, config: RunnableConfig | None = None
-    ) -> Configuration:
-        """Create a Configuration instance from a RunnableConfig object."""
-        configurable = (config.get("configurable") or {}) if config else {}
-        _fields = {f.name for f in fields(cls) if f.init}
-        return cls(**{k: v for k, v in configurable.items() if k in _fields})

surfsense_backend/app/agents/researcher/qna_agent/default_prompts.py

@@ -1,201 +0,0 @@
-"""Default system prompts for Q&A agent.
-
-The prompt system is modular with 3 parts:
-- Part 1 (Base): Core instructions for answering questions (no citations)
-- Part 2 (Citations): Citation-specific instructions and formatting rules
-- Part 3 (Custom): User's custom instructions (empty by default)
-
-Combinations:
-- Part 1 only: Answers without citations
-- Part 1 + Part 2: Answers with citations
-- Part 1 + Part 2 + Part 3: Answers with citations and custom instructions
-"""
-
-# Part 1: Base system prompt for answering without citations
-DEFAULT_QNA_BASE_PROMPT = """Today's date: {date}
-You are SurfSense, an advanced AI research assistant that provides detailed, well-researched answers to user questions by synthesizing information from multiple personal knowledge sources.{language_instruction}
-{chat_history_section}
-<knowledge_sources>
-- EXTENSION: "Web content saved via SurfSense browser extension" (personal browsing history)
-- FILE: "User-uploaded documents (PDFs, Word, etc.)" (personal files)
-- SLACK_CONNECTOR: "Slack conversations and shared content" (personal workspace communications)
-- NOTION_CONNECTOR: "Notion workspace pages and databases" (personal knowledge management)
-- YOUTUBE_VIDEO: "YouTube video transcripts and metadata" (personally saved videos)
-- GITHUB_CONNECTOR: "GitHub repository content and issues" (personal repositories and interactions)
-- ELASTICSEARCH_CONNECTOR: "Elasticsearch indexed documents and data" (personal Elasticsearch instances and custom data sources)
-- LINEAR_CONNECTOR: "Linear project issues and discussions" (personal project management)
-- JIRA_CONNECTOR: "Jira project issues, tickets, and comments" (personal project tracking)
-- CONFLUENCE_CONNECTOR: "Confluence pages and comments" (personal project documentation)
-- CLICKUP_CONNECTOR: "ClickUp tasks and project data" (personal task management)
-- GOOGLE_CALENDAR_CONNECTOR: "Google Calendar events, meetings, and schedules" (personal calendar and time management)
-- GOOGLE_GMAIL_CONNECTOR: "Google Gmail emails and conversations" (personal emails and communications)
-- DISCORD_CONNECTOR: "Discord server conversations and shared content" (personal community communications)
-- AIRTABLE_CONNECTOR: "Airtable records, tables, and database content" (personal data management and organization)
-- TAVILY_API: "Tavily search API results" (personalized search results)
-- LINKUP_API: "Linkup search API results" (personalized search results)
-- LUMA_CONNECTOR: "Luma events"
-- WEBCRAWLER_CONNECTOR: "Webpages indexed by SurfSense" (personally selected websites)
-</knowledge_sources>
-
-<instructions>
-1. Review the chat history to understand the conversation context and any previous topics discussed.
-2. Carefully analyze all provided documents in the <document> sections.
-3. Extract relevant information that directly addresses the user's question.
-4. Provide a comprehensive, detailed answer using information from the user's personal knowledge sources.
-5. Structure your answer logically and conversationally, as if having a detailed discussion with the user.
-6. Use your own words to synthesize and connect ideas from the documents.
-7. If documents contain conflicting information, acknowledge this and present both perspectives.
-8. If the user's question cannot be fully answered with the provided documents, clearly state what information is missing.
-9. Provide actionable insights and practical information when relevant to the user's question.
-10. Use the chat history to maintain conversation continuity and refer to previous discussions when relevant.
-11. Remember that all knowledge sources contain personal information - provide answers that reflect this personal context.
-12. Be conversational and engaging while maintaining accuracy.
-</instructions>
-
-<format>
-- Write in a clear, conversational tone suitable for detailed Q&A discussions
-- Provide comprehensive answers that thoroughly address the user's question
-- Use appropriate paragraphs and structure for readability
-- ALWAYS provide personalized answers that reflect the user's own knowledge and context
-- Be thorough and detailed in your explanations while remaining focused on the user's specific question
-- If asking follow-up questions would be helpful, suggest them at the end of your response
-</format>
-
-<user_query_instructions>
-When you see a user query, focus exclusively on providing a detailed, comprehensive answer using information from the provided documents, which contain the user's personal knowledge and data.
-
-Make sure your response:
-1. Considers the chat history for context and conversation continuity
-2. Directly and thoroughly answers the user's question with personalized information from their own knowledge sources
-3. Is conversational, engaging, and detailed
-4. Acknowledges the personal nature of the information being provided
-5. Offers follow-up suggestions when appropriate
-</user_query_instructions>
-"""
-
-# Part 2: Citation-specific instructions to add citation capabilities
-DEFAULT_QNA_CITATION_INSTRUCTIONS = """
-<citation_instructions>
-CRITICAL CITATION REQUIREMENTS:
-
-1. For EVERY piece of information you include from the documents, add a citation in the format [citation:chunk_id] where chunk_id is the exact value from the `<chunk id='...'>` tag inside `<document_content>`.
-2. Make sure ALL factual statements from the documents have proper citations.
-3. If multiple chunks support the same point, include all relevant citations [citation:chunk_id1], [citation:chunk_id2].
-4. You MUST use the exact chunk_id values from the `<chunk id='...'>` attributes. Do not create your own citation numbers.
-5. Every citation MUST be in the format [citation:chunk_id] where chunk_id is the exact chunk id value.
-6. Never modify or change the chunk_id - always use the original values exactly as provided in the chunk tags.
-7. Do not return citations as clickable links.
-8. Never format citations as markdown links like "([citation:5](https://example.com))". Always use plain square brackets only.
-9. Citations must ONLY appear as [citation:chunk_id] or [citation:chunk_id1], [citation:chunk_id2] format - never with parentheses, hyperlinks, or other formatting.
-10. Never make up chunk IDs. Only use chunk_id values that are explicitly provided in the `<chunk id='...'>` tags.
-11. If you are unsure about a chunk_id, do not include a citation rather than guessing or making one up.
-
-<document_structure_example>
-The documents you receive are structured like this:
-
-<document>
-<document_metadata>
-  <document_id>42</document_id>
-  <document_type>GITHUB_CONNECTOR</document_type>
-  <title><![CDATA[Some repo / file / issue title]]></title>
-  <url><![CDATA[https://example.com]]></url>
-  <metadata_json><![CDATA[{{"any":"other metadata"}}]]></metadata_json>
-</document_metadata>
-
-<document_content>
-  <chunk id='123'><![CDATA[First chunk text...]]></chunk>
-  <chunk id='124'><![CDATA[Second chunk text...]]></chunk>
-</document_content>
-</document>
-
-IMPORTANT: You MUST cite using the chunk ids (e.g. 123, 124). Do NOT cite document_id.
-</document_structure_example>
-
-<citation_format>
-- Every fact from the documents must have a citation in the format [citation:chunk_id] where chunk_id is the EXACT id value from a `<chunk id='...'>` tag
-- Citations should appear at the end of the sentence containing the information they support
-- Multiple citations should be separated by commas: [citation:chunk_id1], [citation:chunk_id2], [citation:chunk_id3]
-- No need to return references section. Just citations in answer.
-- NEVER create your own citation format - use the exact chunk_id values from the documents in the [citation:chunk_id] format
-- NEVER format citations as clickable links or as markdown links like "([citation:5](https://example.com))". Always use plain square brackets only
-- NEVER make up chunk IDs if you are unsure about the chunk_id. It is better to omit the citation than to guess
-</citation_format>
-
-<citation_examples>
-CORRECT citation formats:
-- [citation:5]
-- [citation:chunk_id1], [citation:chunk_id2], [citation:chunk_id3]
-
-INCORRECT citation formats (DO NOT use):
-- Using parentheses and markdown links: ([citation:5](https://github.com/MODSetter/SurfSense))
-- Using parentheses around brackets: ([citation:5])
-- Using hyperlinked text: [link to source 5](https://example.com)
-- Using footnote style: ... library¹
-- Making up source IDs when source_id is unknown
-- Using old IEEE format: [1], [2], [3]
-- Using source types instead of IDs: [citation:GITHUB_CONNECTOR] instead of [citation:5]
-</citation_examples>
-
-<citation_output_example>
-Based on your GitHub repositories and video content, Python's asyncio library provides tools for writing concurrent code using the async/await syntax [citation:5]. It's particularly useful for I/O-bound and high-level structured network code [citation:5].
-
-The key advantage of asyncio is that it can improve performance by allowing other code to run while waiting for I/O operations to complete [citation:12]. This makes it excellent for scenarios like web scraping, API calls, database operations, or any situation where your program spends time waiting for external resources.
-
-However, from your video learning, it's important to note that asyncio is not suitable for CPU-bound tasks as it runs on a single thread [citation:12]. For computationally intensive work, you'd want to use multiprocessing instead.
-</citation_output_example>
-</citation_instructions>
-"""
-
-# Part 3: User's custom instructions (empty by default, can be set by user from UI)
-DEFAULT_QNA_CUSTOM_INSTRUCTIONS = ""
-
-# Full prompt with all parts combined (for backward compatibility and migration)
-DEFAULT_QNA_CITATION_PROMPT = (
-    DEFAULT_QNA_BASE_PROMPT
-    + DEFAULT_QNA_CITATION_INSTRUCTIONS
-    + DEFAULT_QNA_CUSTOM_INSTRUCTIONS
-)
-
-DEFAULT_QNA_NO_DOCUMENTS_PROMPT = """Today's date: {date}
-You are SurfSense, an advanced AI research assistant that provides helpful, detailed answers to user questions in a conversational manner.{language_instruction}
-{chat_history_section}
-<context>
-The user has asked a question but there are no specific documents from their personal knowledge base available to answer it. You should provide a helpful response based on:
-1. The conversation history and context
-2. Your general knowledge and expertise
-3. Understanding of the user's needs and interests based on our conversation
-</context>
-
-<instructions>
-1. Provide a comprehensive, helpful answer to the user's question
-2. Draw upon the conversation history to understand context and the user's specific needs
-3. Use your general knowledge to provide accurate, detailed information
-4. Be conversational and engaging, as if having a detailed discussion with the user
-5. Acknowledge when you're drawing from general knowledge rather than their personal sources
-6. Provide actionable insights and practical information when relevant
-7. Structure your answer logically and clearly
-8. If the question would benefit from personalized information from their knowledge base, gently suggest they might want to add relevant content to SurfSense
-9. Be honest about limitations while still being maximally helpful
-10. Maintain the helpful, knowledgeable tone that users expect from SurfSense
-</instructions>
-
-<format>
-- Write in a clear, conversational tone suitable for detailed Q&A discussions
-- Provide comprehensive answers that thoroughly address the user's question
-- Use appropriate paragraphs and structure for readability
-- No citations are needed since you're using general knowledge
-- Be thorough and detailed in your explanations while remaining focused on the user's specific question
-- If asking follow-up questions would be helpful, suggest them at the end of your response
-- When appropriate, mention that adding relevant content to their SurfSense knowledge base could provide more personalized answers
-</format>
-
-<user_query_instructions>
-When answering the user's question without access to their personal documents:
-1. Review the chat history to understand conversation context and maintain continuity
-2. Provide the most helpful and comprehensive answer possible using general knowledge
-3. Be conversational and engaging
-4. Draw upon conversation history for context
-5. Be clear that you're providing general information
-6. Suggest ways the user could get more personalized answers by expanding their knowledge base when relevant
-</user_query_instructions>
-"""

surfsense_backend/app/agents/researcher/qna_agent/graph.py

@@ -1,21 +0,0 @@
-from langgraph.graph import StateGraph
-
-from .configuration import Configuration
-from .nodes import answer_question, rerank_documents
-from .state import State
-
-# Define a new graph
-workflow = StateGraph(State, config_schema=Configuration)
-
-# Add the nodes to the graph
-workflow.add_node("rerank_documents", rerank_documents)
-workflow.add_node("answer_question", answer_question)
-
-# Connect the nodes
-workflow.add_edge("__start__", "rerank_documents")
-workflow.add_edge("rerank_documents", "answer_question")
-workflow.add_edge("answer_question", "__end__")
-
-# Compile the workflow into an executable graph
-graph = workflow.compile()
-graph.name = "SurfSense QnA Agent"  # This defines the custom name in LangSmith

surfsense_backend/app/agents/researcher/qna_agent/nodes.py

@@ -1,297 +0,0 @@
-import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage, SystemMessage
-from langchain_core.runnables import RunnableConfig
-from langgraph.types import StreamWriter
-from sqlalchemy import select
-
-from app.db import SearchSpace
-from app.services.reranker_service import RerankerService
-
-from ..utils import (
-    calculate_token_count,
-    format_documents_section,
-    langchain_chat_history_to_str,
-    optimize_documents_for_token_limit,
-)
-from .configuration import Configuration
-from .default_prompts import (
-    DEFAULT_QNA_BASE_PROMPT,
-    DEFAULT_QNA_CITATION_INSTRUCTIONS,
-    DEFAULT_QNA_NO_DOCUMENTS_PROMPT,
-)
-from .state import State
-
-
-def _build_language_instruction(language: str | None = None):
-    """Build language instruction for prompts."""
-    if language:
-        return f"\n\nIMPORTANT: Please respond in {language} language. All your responses, explanations, and analysis should be written in {language}."
-    return ""
-
-
-def _build_chat_history_section(chat_history: str | None = None):
-    """Build chat history section for prompts."""
-    if chat_history:
-        return f"""
-<chat_history>
-{chat_history if chat_history else "NO CHAT HISTORY PROVIDED"}
-</chat_history>
-"""
-    return """
-<chat_history>
-NO CHAT HISTORY PROVIDED
-</chat_history>
-"""
-
-
-def _format_system_prompt(
-    prompt_template: str,
-    chat_history: str | None = None,
-    language: str | None = None,
-):
-    """Format a system prompt template with dynamic values."""
-    date = datetime.datetime.now().strftime("%Y-%m-%d")
-    language_instruction = _build_language_instruction(language)
-    chat_history_section = _build_chat_history_section(chat_history)
-
-    return prompt_template.format(
-        date=date,
-        language_instruction=language_instruction,
-        chat_history_section=chat_history_section,
-    )
-
-
-async def rerank_documents(state: State, config: RunnableConfig) -> dict[str, Any]:
-    """
-    Rerank the documents based on relevance to the user's question.
-
-    This node takes the relevant documents provided in the configuration,
-    reranks them using the reranker service based on the user's query,
-    and updates the state with the reranked documents.
-
-    Documents are now document-grouped with a `chunks` list. Reranking is done
-    using the concatenated `content` field, and the full structure (including
-    `chunks`) is preserved for proper citation formatting.
-
-    If reranking is disabled, returns the original documents without processing.
-
-    Returns:
-        Dict containing the reranked documents.
-    """
-    # Get configuration and relevant documents
-    configuration = Configuration.from_runnable_config(config)
-    documents = configuration.relevant_documents
-    user_query = configuration.user_query
-    reformulated_query = configuration.reformulated_query
-
-    # If no documents were provided, return empty list
-    if not documents or len(documents) == 0:
-        return {"reranked_documents": []}
-
-    # Get reranker service from app config
-    reranker_service = RerankerService.get_reranker_instance()
-
-    # If reranking is not enabled, sort by existing score and return
-    if not reranker_service:
-        print("Reranking is disabled. Sorting documents by existing score.")
-        sorted_documents = sorted(
-            documents, key=lambda x: x.get("score", 0), reverse=True
-        )
-        return {"reranked_documents": sorted_documents}
-
-    # Perform reranking
-    try:
-        # Pass documents directly to reranker - it will use:
-        # - "content" (concatenated chunk text) for scoring
-        # - "chunk_id" (primary chunk id) for matching
-        # The full document structure including "chunks" is preserved
-        reranked_docs = reranker_service.rerank_documents(
-            user_query + "\n" + reformulated_query, documents
-        )
-
-        # Sort by score in descending order
-        reranked_docs.sort(key=lambda x: x.get("score", 0), reverse=True)
-
-        print(f"Reranked {len(reranked_docs)} documents for Q&A query: {user_query}")
-
-        return {"reranked_documents": reranked_docs}
-
-    except Exception as e:
-        print(f"Error during reranking: {e!s}")
-        # Fall back to original documents if reranking fails
-        return {"reranked_documents": documents}
-
-
-async def answer_question(
-    state: State, config: RunnableConfig, writer: StreamWriter
-) -> dict[str, Any]:
-    """
-    Answer the user's question using the provided documents with real-time streaming.
-
-    This node takes the relevant documents provided in the configuration and uses
-    an LLM to generate a comprehensive answer to the user's question with
-    proper citations. The citations follow [citation:chunk_id] format using chunk IDs from the
-    `<chunk id='...'>` tags in the provided documents. If no documents are provided, it will use chat history to generate
-    an answer.
-
-    The response is streamed token-by-token for real-time updates to the frontend.
-
-    Returns:
-        Dict containing the final answer in the "final_answer" key.
-    """
-    from app.services.llm_service import get_fast_llm
-
-    # Get configuration and relevant documents from configuration
-    configuration = Configuration.from_runnable_config(config)
-    documents = state.reranked_documents
-    user_query = configuration.user_query
-    search_space_id = configuration.search_space_id
-    language = configuration.language
-
-    # Get streaming service from state
-    streaming_service = state.streaming_service
-
-    # Fetch search space to get QnA configuration
-    result = await state.db_session.execute(
-        select(SearchSpace).where(SearchSpace.id == search_space_id)
-    )
-    search_space = result.scalar_one_or_none()
-
-    if not search_space:
-        error_message = f"Search space {search_space_id} not found"
-        print(error_message)
-        raise RuntimeError(error_message)
-
-    # Get QnA configuration from search space
-    citations_enabled = search_space.citations_enabled
-    custom_instructions_text = search_space.qna_custom_instructions or ""
-
-    # Use constants for base prompt and citation instructions
-    qna_base_prompt = DEFAULT_QNA_BASE_PROMPT
-    qna_citation_instructions = (
-        DEFAULT_QNA_CITATION_INSTRUCTIONS if citations_enabled else ""
-    )
-    qna_custom_instructions = (
-        f"\n<special_important_custom_instructions>\n{custom_instructions_text}\n</special_important_custom_instructions>"
-        if custom_instructions_text
-        else ""
-    )
-
-    # Get search space's fast LLM
-    llm = await get_fast_llm(state.db_session, search_space_id)
-    if not llm:
-        error_message = f"No fast LLM configured for search space {search_space_id}"
-        print(error_message)
-        raise RuntimeError(error_message)
-
-    # Determine if we have documents and optimize for token limits
-    has_documents_initially = documents and len(documents) > 0
-    chat_history_str = langchain_chat_history_to_str(state.chat_history)
-
-    if has_documents_initially:
-        # Compose the full citation prompt: base + citation instructions + custom instructions
-        full_citation_prompt_template = (
-            qna_base_prompt + qna_citation_instructions + qna_custom_instructions
-        )
-
-        # Create base message template for token calculation (without documents)
-        base_human_message_template = f"""
-        
-        User's question:
-        <user_query>
-            {user_query}
-        </user_query>
-        
-        Please provide a detailed, comprehensive answer to the user's question using the information from their personal knowledge sources. Make sure to cite all information appropriately and engage in a conversational manner.
-        """
-
-        # Use initial system prompt for token calculation
-        initial_system_prompt = _format_system_prompt(
-            full_citation_prompt_template, chat_history_str, language
-        )
-        base_messages = [
-            SystemMessage(content=initial_system_prompt),
-            HumanMessage(content=base_human_message_template),
-        ]
-
-        # Optimize documents to fit within token limits
-        optimized_documents, has_optimized_documents = (
-            optimize_documents_for_token_limit(documents, base_messages, llm.model)
-        )
-
-        # Update state based on optimization result
-        documents = optimized_documents
-        has_documents = has_optimized_documents
-    else:
-        has_documents = False
-
-    # Choose system prompt based on final document availability
-    # With documents: use base + citation instructions + custom instructions
-    # Without documents: use the default no-documents prompt from constants
-    if has_documents:
-        full_citation_prompt_template = (
-            qna_base_prompt + qna_citation_instructions + qna_custom_instructions
-        )
-        system_prompt = _format_system_prompt(
-            full_citation_prompt_template, chat_history_str, language
-        )
-    else:
-        system_prompt = _format_system_prompt(
-            DEFAULT_QNA_NO_DOCUMENTS_PROMPT + qna_custom_instructions,
-            chat_history_str,
-            language,
-        )
-
-    # Generate documents section
-    documents_text = (
-        format_documents_section(
-            documents, "Source material from your personal knowledge base"
-        )
-        if has_documents
-        else ""
-    )
-
-    # Create final human message content
-    instruction_text = (
-        "Please provide a detailed, comprehensive answer to the user's question using the information from their personal knowledge sources. Make sure to cite all information appropriately and engage in a conversational manner."
-        if has_documents
-        else "Please provide a helpful answer to the user's question based on our conversation history and your general knowledge. Engage in a conversational manner."
-    )
-
-    human_message_content = f"""
-    {documents_text}
-    
-    User's question:
-    <user_query>
-        {user_query}
-    </user_query>
-    
-    {instruction_text}
-    """
-
-    # Create final messages for the LLM
-    messages_with_chat_history = [
-        SystemMessage(content=system_prompt),
-        HumanMessage(content=human_message_content),
-    ]
-
-    # Log final token count
-    total_tokens = calculate_token_count(messages_with_chat_history, llm.model)
-    print(f"Final token count: {total_tokens}")
-
-    # Stream the LLM response token by token
-    final_answer = ""
-
-    async for chunk in llm.astream(messages_with_chat_history):
-        # Extract the content from the chunk
-        if hasattr(chunk, "content") and chunk.content:
-            token = chunk.content
-            final_answer += token
-
-            # Stream the token to the frontend via custom stream
-            if streaming_service:
-                writer({"yield_value": streaming_service.format_text_chunk(token)})
-
-    return {"final_answer": final_answer}

surfsense_backend/app/agents/researcher/qna_agent/state.py

@@ -1,32 +0,0 @@
-"""Define the state structures for the agent."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import Any
-
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from app.services.streaming_service import StreamingService
-
-
-@dataclass
-class State:
-    """Defines the dynamic state for the Q&A agent during execution.
-
-    This state tracks the database session, chat history, and the outputs
-    generated by the agent's nodes during question answering.
-    See: https://langchain-ai.github.io/langgraph/concepts/low_level/#state
-    for more information.
-    """
-
-    # Runtime context
-    db_session: AsyncSession
-
-    # Streaming service for real-time token streaming
-    streaming_service: StreamingService | None = None
-
-    chat_history: list[Any] | None = field(default_factory=list)
-    # OUTPUT: Populated by agent nodes
-    reranked_documents: list[Any] | None = None
-    final_answer: str | None = None

surfsense_backend/app/agents/researcher/state.py

@@ -1,38 +0,0 @@
-"""Define the state structures for the agent."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import Any
-
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from app.services.streaming_service import StreamingService
-
-
-@dataclass
-class State:
-    """Defines the dynamic state for the agent during execution.
-
-    This state tracks the database session and the outputs generated by the agent's nodes.
-    See: https://langchain-ai.github.io/langgraph/concepts/low_level/#state
-    for more information.
-    """
-
-    # Runtime context (not part of actual graph state)
-    db_session: AsyncSession
-
-    # Streaming service
-    streaming_service: StreamingService
-
-    chat_history: list[Any] | None = field(default_factory=list)
-
-    reformulated_query: str | None = field(default=None)
-    further_questions: Any | None = field(default=None)
-
-    # Temporary field to hold reranked documents from sub-agents for further question generation
-    reranked_documents: list[Any] | None = field(default=None)
-
-    # OUTPUT: Populated by agent nodes
-    # Using field to explicitly mark as part of state
-    final_written_report: str | None = field(default=None)

surfsense_backend/app/agents/researcher/utils.py

@@ -1,292 +0,0 @@
-import json
-from typing import Any, NamedTuple
-
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
-from litellm import get_model_info, token_counter
-
-
-class DocumentTokenInfo(NamedTuple):
-    """Information about a document and its token cost."""
-
-    index: int
-    document: dict[str, Any]
-    formatted_content: str
-    token_count: int
-
-
-def get_connector_emoji(connector_name: str) -> str:
-    """Get an appropriate emoji for a connector type."""
-    connector_emojis = {
-        "YOUTUBE_VIDEO": "📹",
-        "EXTENSION": "🧩",
-        "FILE": "📄",
-        "SLACK_CONNECTOR": "💬",
-        "NOTION_CONNECTOR": "📘",
-        "GITHUB_CONNECTOR": "🐙",
-        "LINEAR_CONNECTOR": "📊",
-        "JIRA_CONNECTOR": "🎫",
-        "DISCORD_CONNECTOR": "🗨️",
-        "TAVILY_API": "🔍",
-        "LINKUP_API": "🔗",
-        "BAIDU_SEARCH_API": "🇨🇳",
-        "GOOGLE_CALENDAR_CONNECTOR": "📅",
-        "AIRTABLE_CONNECTOR": "🗃️",
-        "LUMA_CONNECTOR": "✨",
-        "ELASTICSEARCH_CONNECTOR": "⚡",
-        "WEBCRAWLER_CONNECTOR": "🌐",
-        "BOOKSTACK_CONNECTOR": "📚",
-        "NOTE": "📝",
-    }
-    return connector_emojis.get(connector_name, "🔎")
-
-
-def get_connector_friendly_name(connector_name: str) -> str:
-    """Convert technical connector IDs to user-friendly names."""
-    connector_friendly_names = {
-        "YOUTUBE_VIDEO": "YouTube",
-        "EXTENSION": "Browser Extension",
-        "FILE": "Files",
-        "SLACK_CONNECTOR": "Slack",
-        "NOTION_CONNECTOR": "Notion",
-        "GITHUB_CONNECTOR": "GitHub",
-        "LINEAR_CONNECTOR": "Linear",
-        "JIRA_CONNECTOR": "Jira",
-        "CONFLUENCE_CONNECTOR": "Confluence",
-        "GOOGLE_CALENDAR_CONNECTOR": "Google Calendar",
-        "DISCORD_CONNECTOR": "Discord",
-        "TAVILY_API": "Tavily Search",
-        "LINKUP_API": "Linkup Search",
-        "BAIDU_SEARCH_API": "Baidu Search",
-        "AIRTABLE_CONNECTOR": "Airtable",
-        "LUMA_CONNECTOR": "Luma",
-        "ELASTICSEARCH_CONNECTOR": "Elasticsearch",
-        "WEBCRAWLER_CONNECTOR": "Web Pages",
-        "BOOKSTACK_CONNECTOR": "BookStack",
-        "NOTE": "Notes",
-    }
-    return connector_friendly_names.get(connector_name, connector_name)
-
-
-def convert_langchain_messages_to_dict(
-    messages: list[BaseMessage],
-) -> list[dict[str, str]]:
-    """Convert LangChain messages to format expected by token_counter."""
-    role_mapping = {"system": "system", "human": "user", "ai": "assistant"}
-
-    converted_messages = []
-    for msg in messages:
-        role = role_mapping.get(getattr(msg, "type", None), "user")
-        converted_messages.append({"role": role, "content": str(msg.content)})
-
-    return converted_messages
-
-
-def format_document_for_citation(document: dict[str, Any]) -> str:
-    """Format a single document for citation in the new document+chunks XML format.
-
-    IMPORTANT:
-    - Citations must reference real DB chunk IDs: `[citation:<chunk_id>]`
-    - Document metadata is included under <document_metadata>, but citations are NOT document_id-based.
-    """
-
-    def _to_cdata(value: Any) -> str:
-        text = "" if value is None else str(value)
-        # Safely nest CDATA even if the content includes "]]>"
-        return "<![CDATA[" + text.replace("]]>", "]]]]><![CDATA[>") + "]]>"
-
-    doc_info = document.get("document", {}) or {}
-    metadata = doc_info.get("metadata", {}) or {}
-
-    doc_id = doc_info.get("id", "")
-    title = doc_info.get("title", "")
-    document_type = doc_info.get("document_type", "CRAWLED_URL")
-    url = (
-        metadata.get("url")
-        or metadata.get("source")
-        or metadata.get("page_url")
-        or metadata.get("VisitedWebPageURL")
-        or ""
-    )
-
-    metadata_json = json.dumps(metadata, ensure_ascii=False)
-
-    chunks = document.get("chunks") or []
-    if not chunks:
-        # Fallback: treat `content` as a single chunk (no chunk_id available for citation)
-        chunks = [{"chunk_id": "", "content": document.get("content", "")}]
-
-    chunks_xml = "\n".join(
-        [
-            f"<chunk id='{chunk.get('chunk_id', '')}'>{_to_cdata(chunk.get('content', ''))}</chunk>"
-            for chunk in chunks
-        ]
-    )
-
-    return f"""<document>
-<document_metadata>
-<document_id>{doc_id}</document_id>
-<document_type>{document_type}</document_type>
-<title>{_to_cdata(title)}</title>
-<url>{_to_cdata(url)}</url>
-<metadata_json>{_to_cdata(metadata_json)}</metadata_json>
-</document_metadata>
-
-<document_content>
-{chunks_xml}
-</document_content>
-</document>"""
-
-
-def format_documents_section(
-    documents: list[dict[str, Any]], section_title: str = "Source material"
-) -> str:
-    """Format multiple documents into a complete documents section."""
-    if not documents:
-        return ""
-
-    formatted_docs = [format_document_for_citation(doc) for doc in documents]
-
-    return f"""{section_title}:
-    <documents>
-    {chr(10).join(formatted_docs)}
-    </documents>"""
-
-
-def calculate_document_token_costs(
-    documents: list[dict[str, Any]], model: str
-) -> list[DocumentTokenInfo]:
-    """Pre-calculate token costs for each document."""
-    document_token_info = []
-
-    for i, doc in enumerate(documents):
-        formatted_doc = format_document_for_citation(doc)
-
-        # Calculate token count for this document
-        token_count = token_counter(
-            messages=[{"role": "user", "content": formatted_doc}], model=model
-        )
-
-        document_token_info.append(
-            DocumentTokenInfo(
-                index=i,
-                document=doc,
-                formatted_content=formatted_doc,
-                token_count=token_count,
-            )
-        )
-
-    return document_token_info
-
-
-def find_optimal_documents_with_binary_search(
-    document_tokens: list[DocumentTokenInfo], available_tokens: int
-) -> list[DocumentTokenInfo]:
-    """Use binary search to find the maximum number of documents that fit within token limit."""
-    if not document_tokens or available_tokens <= 0:
-        return []
-
-    left, right = 0, len(document_tokens)
-    optimal_docs = []
-
-    while left <= right:
-        mid = (left + right) // 2
-        current_docs = document_tokens[:mid]
-        current_token_sum = sum(doc_info.token_count for doc_info in current_docs)
-
-        if current_token_sum <= available_tokens:
-            optimal_docs = current_docs
-            left = mid + 1
-        else:
-            right = mid - 1
-
-    return optimal_docs
-
-
-def get_model_context_window(model_name: str) -> int:
-    """Get the total context window size for a model (input + output tokens)."""
-    try:
-        model_info = get_model_info(model_name)
-        context_window = model_info.get("max_input_tokens", 4096)  # Default fallback
-        return context_window
-    except Exception as e:
-        print(
-            f"Warning: Could not get model info for {model_name}, using default 4096 tokens. Error: {e}"
-        )
-        return 4096  # Conservative fallback
-
-
-def optimize_documents_for_token_limit(
-    documents: list[dict[str, Any]], base_messages: list[BaseMessage], model_name: str
-) -> tuple[list[dict[str, Any]], bool]:
-    """
-    Optimize documents to fit within token limits using binary search.
-
-    Args:
-        documents: List of documents with content and metadata
-        base_messages: Base messages without documents (chat history + system + human message template)
-        model_name: Model name for token counting (required)
-        output_token_buffer: Number of tokens to reserve for model output
-
-    Returns:
-        Tuple of (optimized_documents, has_documents_remaining)
-    """
-    if not documents:
-        return [], False
-
-    model = model_name
-    context_window = get_model_context_window(model)
-
-    # Calculate base token cost
-    base_messages_dict = convert_langchain_messages_to_dict(base_messages)
-    base_tokens = token_counter(messages=base_messages_dict, model=model)
-    available_tokens_for_docs = context_window - base_tokens
-
-    print(
-        f"Token optimization: Context window={context_window}, Base={base_tokens}, Available for docs={available_tokens_for_docs}"
-    )
-
-    if available_tokens_for_docs <= 0:
-        print("No tokens available for documents after base content and output buffer")
-        return [], False
-
-    # Calculate token costs for all documents
-    document_token_info = calculate_document_token_costs(documents, model)
-
-    # Find optimal number of documents using binary search
-    optimal_doc_info = find_optimal_documents_with_binary_search(
-        document_token_info, available_tokens_for_docs
-    )
-
-    # Extract the original document objects
-    optimized_documents = [doc_info.document for doc_info in optimal_doc_info]
-    has_documents_remaining = len(optimized_documents) > 0
-
-    print(
-        f"Token optimization result: Using {len(optimized_documents)}/{len(documents)} documents"
-    )
-
-    return optimized_documents, has_documents_remaining
-
-
-def calculate_token_count(messages: list[BaseMessage], model_name: str) -> int:
-    """Calculate token count for a list of LangChain messages."""
-    model = model_name
-    messages_dict = convert_langchain_messages_to_dict(messages)
-    return token_counter(messages=messages_dict, model=model)
-
-
-def langchain_chat_history_to_str(chat_history: list[BaseMessage]) -> str:
-    """
-    Convert a list of chat history messages to a string.
-    """
-    chat_history_str = ""
-
-    for chat_message in chat_history:
-        if isinstance(chat_message, HumanMessage):
-            chat_history_str += f"<user>{chat_message.content}</user>\n"
-        elif isinstance(chat_message, AIMessage):
-            chat_history_str += f"<assistant>{chat_message.content}</assistant>\n"
-        elif isinstance(chat_message, SystemMessage):
-            chat_history_str += f"<system>{chat_message.content}</system>\n"
-
-    return chat_history_str

surfsense_backend/app/db.py

@@ -77,10 +77,6 @@ class SearchSourceConnectorType(str, Enum):
     BOOKSTACK_CONNECTOR = "BOOKSTACK_CONNECTOR"
 
 
-class ChatType(str, Enum):
-    QNA = "QNA"
-
-
 class LiteLLMProvider(str, Enum):
     """
     Enum for LLM providers supported by LiteLLM.
@@ -317,21 +313,6 @@ class BaseModel(Base):
     id = Column(Integer, primary_key=True, index=True)
 
 
-class Chat(BaseModel, TimestampMixin):
-    __tablename__ = "chats"
-
-    type = Column(SQLAlchemyEnum(ChatType), nullable=False)
-    title = Column(String, nullable=False, index=True)
-    initial_connectors = Column(ARRAY(String), nullable=True)
-    messages = Column(JSON, nullable=False)
-    state_version = Column(BigInteger, nullable=False, default=1)
-
-    search_space_id = Column(
-        Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
-    )
-    search_space = relationship("SearchSpace", back_populates="chats")
-
-
 class NewChatMessageRole(str, Enum):
     """Role enum for new chat messages."""
 
@@ -363,9 +344,6 @@ class NewChatThread(BaseModel, TimestampMixin):
     search_space_id = Column(
         Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
     )
-    user_id = Column(
-        UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False
-    )
 
     # Relationships
     search_space = relationship("SearchSpace", back_populates="new_chat_threads")
@@ -445,23 +423,6 @@ class Chunk(BaseModel, TimestampMixin):
     document = relationship("Document", back_populates="chunks")
 
 
-class Podcast(BaseModel, TimestampMixin):
-    __tablename__ = "podcasts"
-
-    title = Column(String, nullable=False, index=True)
-    podcast_transcript = Column(JSON, nullable=False, default={})
-    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(
-        Integer, ForeignKey("searchspaces.id", ondelete="CASCADE"), nullable=False
-    )
-    search_space = relationship("SearchSpace", back_populates="podcasts")
-
-
 class SearchSpace(BaseModel, TimestampMixin):
     __tablename__ = "searchspaces"
 
@@ -492,18 +453,6 @@ class SearchSpace(BaseModel, TimestampMixin):
         order_by="Document.id",
         cascade="all, delete-orphan",
     )
-    podcasts = relationship(
-        "Podcast",
-        back_populates="search_space",
-        order_by="Podcast.id",
-        cascade="all, delete-orphan",
-    )
-    chats = relationship(
-        "Chat",
-        back_populates="search_space",
-        order_by="Chat.id",
-        cascade="all, delete-orphan",
-    )
     new_chat_threads = relationship(
         "NewChatThread",
         back_populates="search_space",

surfsense_backend/app/routes/__init__.py

@@ -3,7 +3,6 @@ from fastapi import APIRouter
 from .airtable_add_connector_route import (
     router as airtable_add_connector_router,
 )
-from .chats_routes import router as chats_router
 from .documents_routes import router as documents_router
 from .editor_routes import router as editor_router
 from .google_calendar_add_connector_route import (
@@ -17,7 +16,6 @@ from .logs_routes import router as logs_router
 from .luma_add_connector_route import router as luma_add_connector_router
 from .new_chat_routes import router as new_chat_router
 from .notes_routes import router as notes_router
-from .podcasts_routes import router as podcasts_router
 from .rbac_routes import router as rbac_router
 from .search_source_connectors_routes import router as search_source_connectors_router
 from .search_spaces_routes import router as search_spaces_router
@@ -29,9 +27,7 @@ router.include_router(rbac_router)  # RBAC routes for roles, members, invites
 router.include_router(editor_router)
 router.include_router(documents_router)
 router.include_router(notes_router)
-router.include_router(podcasts_router)
+router.include_router(new_chat_router)  # Chat with assistant-ui persistence
-router.include_router(chats_router)
-router.include_router(new_chat_router)  # New chat with assistant-ui persistence
 router.include_router(search_source_connectors_router)
 router.include_router(google_calendar_add_connector_router)
 router.include_router(google_gmail_add_connector_router)

surfsense_backend/app/routes/chats_routes.py

@@ -1,617 +0,0 @@
-from fastapi import APIRouter, Depends, HTTPException
-from fastapi.responses import StreamingResponse
-from langchain_core.messages import AIMessage, HumanMessage
-from sqlalchemy.exc import IntegrityError, OperationalError
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.future import select
-from sqlalchemy.orm import selectinload
-
-from app.db import (
-    Chat,
-    Permission,
-    SearchSpace,
-    SearchSpaceMembership,
-    User,
-    get_async_session,
-)
-from app.schemas import (
-    AISDKChatRequest,
-    ChatCreate,
-    ChatRead,
-    ChatReadWithoutMessages,
-    ChatUpdate,
-    NewChatRequest,
-)
-from app.services.new_streaming_service import VercelStreamingService
-from app.tasks.chat.stream_connector_search_results import (
-    stream_connector_search_results,
-)
-from app.tasks.chat.stream_new_chat import stream_new_chat
-from app.users import current_active_user
-from app.utils.rbac import check_permission
-from app.utils.validators import (
-    validate_connectors,
-    validate_document_ids,
-    validate_messages,
-    validate_research_mode,
-    validate_search_space_id,
-    validate_top_k,
-)
-
-router = APIRouter()
-
-
-@router.post("/chat")
-async def handle_chat_data(
-    request: AISDKChatRequest,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    # Validate and sanitize all input data
-    messages = validate_messages(request.messages)
-
-    if messages[-1]["role"] != "user":
-        raise HTTPException(
-            status_code=400, detail="Last message must be a user message"
-        )
-
-    user_query = messages[-1]["content"]
-
-    # Extract and validate data from request
-    request_data = request.data or {}
-    search_space_id = validate_search_space_id(request_data.get("search_space_id"))
-    research_mode = validate_research_mode(request_data.get("research_mode"))
-    selected_connectors = validate_connectors(request_data.get("selected_connectors"))
-    document_ids_to_add_in_context = validate_document_ids(
-        request_data.get("document_ids_to_add_in_context")
-    )
-    top_k = validate_top_k(request_data.get("top_k"))
-    # print("RESQUEST DATA:", request_data)
-    # print("SELECTED CONNECTORS:", selected_connectors)
-
-    # Check if the user has chat access to the search space
-    try:
-        await check_permission(
-            session,
-            user,
-            search_space_id,
-            Permission.CHATS_CREATE.value,
-            "You don't have permission to use chat in this search space",
-        )
-
-        # Get search space with LLM configs (preferences are now stored at search space level)
-        search_space_result = await session.execute(
-            select(SearchSpace)
-            .options(selectinload(SearchSpace.llm_configs))
-            .filter(SearchSpace.id == search_space_id)
-        )
-        search_space = search_space_result.scalars().first()
-
-        language = None
-        llm_configs = []  # Initialize to empty list
-
-        if search_space and search_space.llm_configs:
-            llm_configs = search_space.llm_configs
-
-            # Get language from configured LLM preferences
-            # LLM preferences are now stored on the SearchSpace model
-            from app.config import config as app_config
-
-            for llm_id in [
-                search_space.fast_llm_id,
-                search_space.long_context_llm_id,
-                search_space.strategic_llm_id,
-            ]:
-                if llm_id is not None:
-                    # Check if it's a global config (negative ID)
-                    if llm_id < 0:
-                        # Look in global configs
-                        for global_cfg in app_config.GLOBAL_LLM_CONFIGS:
-                            if global_cfg.get("id") == llm_id:
-                                language = global_cfg.get("language")
-                                if language:
-                                    break
-                    else:
-                        # Look in custom configs
-                        for llm_config in llm_configs:
-                            if llm_config.id == llm_id and getattr(
-                                llm_config, "language", None
-                            ):
-                                language = llm_config.language
-                                break
-                    if language:
-                        break
-
-        if not language and llm_configs:
-            first_llm_config = llm_configs[0]
-            language = getattr(first_llm_config, "language", None)
-
-    except HTTPException:
-        raise HTTPException(
-            status_code=403, detail="You don't have access to this search space"
-        ) from None
-
-    langchain_chat_history = []
-    for message in messages[:-1]:
-        if message["role"] == "user":
-            langchain_chat_history.append(HumanMessage(content=message["content"]))
-        elif message["role"] == "assistant":
-            langchain_chat_history.append(AIMessage(content=message["content"]))
-
-    response = StreamingResponse(
-        stream_connector_search_results(
-            user_query,
-            user.id,
-            search_space_id,
-            session,
-            research_mode,
-            selected_connectors,
-            langchain_chat_history,
-            document_ids_to_add_in_context,
-            language,
-            top_k,
-        )
-    )
-
-    response.headers["x-vercel-ai-data-stream"] = "v1"
-    return response
-
-
-@router.post("/new_chat")
-async def handle_new_chat(
-    request: NewChatRequest,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Handle new chat requests using the SurfSense deep agent.
-
-    This endpoint uses the new deep agent with the Vercel AI SDK
-    Data Stream Protocol (SSE format).
-
-    Args:
-        request: NewChatRequest containing chat_id, user_query, and search_space_id
-        session: Database session
-        user: Current authenticated user
-
-    Returns:
-        StreamingResponse with SSE formatted data
-    """
-    # Validate the user query
-    if not request.user_query or not request.user_query.strip():
-        raise HTTPException(status_code=400, detail="User query cannot be empty")
-
-    # Check if the user has chat access to the search space
-    try:
-        await check_permission(
-            session,
-            user,
-            request.search_space_id,
-            Permission.CHATS_CREATE.value,
-            "You don't have permission to use chat in this search space",
-        )
-    except HTTPException:
-        raise HTTPException(
-            status_code=403, detail="You don't have access to this search space"
-        ) from None
-
-    # Get LLM config ID from search space preferences (optional enhancement)
-    # For now, we use the default global config (-1)
-    llm_config_id = -1
-
-    # Optionally load LLM preferences from search space
-    try:
-        search_space_result = await session.execute(
-            select(SearchSpace).filter(SearchSpace.id == request.search_space_id)
-        )
-        search_space = search_space_result.scalars().first()
-
-        if search_space:
-            # Use strategic_llm_id if available, otherwise fall back to fast_llm_id
-            if search_space.strategic_llm_id is not None:
-                llm_config_id = search_space.strategic_llm_id
-            elif search_space.fast_llm_id is not None:
-                llm_config_id = search_space.fast_llm_id
-    except Exception:
-        # Fall back to default config on any error
-        pass
-
-    # Create the streaming response
-    # chat_id is used as LangGraph's thread_id for automatic chat history management
-    response = StreamingResponse(
-        stream_new_chat(
-            user_query=request.user_query.strip(),
-            user_id=user.id,
-            search_space_id=request.search_space_id,
-            chat_id=request.chat_id,
-            session=session,
-            llm_config_id=llm_config_id,
-            messages=request.messages,  # Pass message history from frontend
-        ),
-        media_type="text/event-stream",
-    )
-
-    # Set the required headers for Vercel AI SDK
-    headers = VercelStreamingService.get_response_headers()
-    for key, value in headers.items():
-        response.headers[key] = value
-
-    return response
-
-
-@router.post("/chats", response_model=ChatRead)
-async def create_chat(
-    chat: ChatCreate,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Create a new chat.
-    Requires CHATS_CREATE permission.
-    """
-    try:
-        await check_permission(
-            session,
-            user,
-            chat.search_space_id,
-            Permission.CHATS_CREATE.value,
-            "You don't have permission to create chats in this search space",
-        )
-        db_chat = Chat(**chat.model_dump())
-        session.add(db_chat)
-        await session.commit()
-        await session.refresh(db_chat)
-        return db_chat
-    except HTTPException:
-        raise
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400,
-            detail="Database constraint violation. Please check your input data.",
-        ) from None
-    except OperationalError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500,
-            detail="An unexpected error occurred while creating the chat.",
-        ) from None
-
-
-@router.get("/chats", response_model=list[ChatReadWithoutMessages])
-async def read_chats(
-    skip: int = 0,
-    limit: int = 100,
-    search_space_id: int | None = None,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    List chats the user has access to.
-    Requires CHATS_READ permission for the search space(s).
-    """
-    # Validate pagination parameters
-    if skip < 0:
-        raise HTTPException(
-            status_code=400, detail="skip must be a non-negative integer"
-        )
-
-    if limit <= 0 or limit > 1000:  # Reasonable upper limit
-        raise HTTPException(status_code=400, detail="limit must be between 1 and 1000")
-
-    # Validate search_space_id if provided
-    if search_space_id is not None and search_space_id <= 0:
-        raise HTTPException(
-            status_code=400, detail="search_space_id must be a positive integer"
-        )
-    try:
-        if search_space_id is not None:
-            # Check permission for specific search space
-            await check_permission(
-                session,
-                user,
-                search_space_id,
-                Permission.CHATS_READ.value,
-                "You don't have permission to read chats in this search space",
-            )
-            # Select specific fields excluding messages
-            query = (
-                select(
-                    Chat.id,
-                    Chat.type,
-                    Chat.title,
-                    Chat.initial_connectors,
-                    Chat.search_space_id,
-                    Chat.created_at,
-                    Chat.state_version,
-                )
-                .filter(Chat.search_space_id == search_space_id)
-                .order_by(Chat.created_at.desc())
-            )
-        else:
-            # Get chats from all search spaces user has membership in
-            query = (
-                select(
-                    Chat.id,
-                    Chat.type,
-                    Chat.title,
-                    Chat.initial_connectors,
-                    Chat.search_space_id,
-                    Chat.created_at,
-                    Chat.state_version,
-                )
-                .join(SearchSpace)
-                .join(SearchSpaceMembership)
-                .filter(SearchSpaceMembership.user_id == user.id)
-                .order_by(Chat.created_at.desc())
-            )
-
-        result = await session.execute(query.offset(skip).limit(limit))
-        return result.all()
-    except HTTPException:
-        raise
-    except OperationalError:
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        raise HTTPException(
-            status_code=500, detail="An unexpected error occurred while fetching chats."
-        ) from None
-
-
-@router.get("/chats/search", response_model=list[ChatReadWithoutMessages])
-async def search_chats(
-    title: str,
-    skip: int = 0,
-    limit: int = 100,
-    search_space_id: int | None = None,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Search chats by title substring.
-    Requires CHATS_READ permission for the search space(s).
-
-    Args:
-        title: Case-insensitive substring to match against chat titles. Required.
-        skip: Number of items to skip from the beginning. Default: 0.
-        limit: Maximum number of items to return. Default: 100.
-        search_space_id: Filter results to a specific search space. Default: None.
-        session: Database session (injected).
-        user: Current authenticated user (injected).
-
-    Returns:
-        List of chats matching the search query.
-
-    Notes:
-        - Title matching uses ILIKE (case-insensitive).
-        - Results are ordered by creation date (most recent first).
-    """
-    # Validate pagination parameters
-    if skip < 0:
-        raise HTTPException(
-            status_code=400, detail="skip must be a non-negative integer"
-        )
-
-    if limit <= 0 or limit > 1000:
-        raise HTTPException(status_code=400, detail="limit must be between 1 and 1000")
-
-    # Validate search_space_id if provided
-    if search_space_id is not None and search_space_id <= 0:
-        raise HTTPException(
-            status_code=400, detail="search_space_id must be a positive integer"
-        )
-
-    try:
-        if search_space_id is not None:
-            # Check permission for specific search space
-            await check_permission(
-                session,
-                user,
-                search_space_id,
-                Permission.CHATS_READ.value,
-                "You don't have permission to read chats in this search space",
-            )
-            # Select specific fields excluding messages
-            query = (
-                select(
-                    Chat.id,
-                    Chat.type,
-                    Chat.title,
-                    Chat.initial_connectors,
-                    Chat.search_space_id,
-                    Chat.created_at,
-                    Chat.state_version,
-                )
-                .filter(Chat.search_space_id == search_space_id)
-                .order_by(Chat.created_at.desc())
-            )
-        else:
-            # Get chats from all search spaces user has membership in
-            query = (
-                select(
-                    Chat.id,
-                    Chat.type,
-                    Chat.title,
-                    Chat.initial_connectors,
-                    Chat.search_space_id,
-                    Chat.created_at,
-                    Chat.state_version,
-                )
-                .join(SearchSpace)
-                .join(SearchSpaceMembership)
-                .filter(SearchSpaceMembership.user_id == user.id)
-                .order_by(Chat.created_at.desc())
-            )
-
-        # Apply title search filter (case-insensitive)
-        query = query.filter(Chat.title.ilike(f"%{title}%"))
-
-        result = await session.execute(query.offset(skip).limit(limit))
-        return result.all()
-    except HTTPException:
-        raise
-    except OperationalError:
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        raise HTTPException(
-            status_code=500,
-            detail="An unexpected error occurred while searching chats.",
-        ) from None
-
-
-@router.get("/chats/{chat_id}", response_model=ChatRead)
-async def read_chat(
-    chat_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Get a specific chat by ID.
-    Requires CHATS_READ permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Chat).filter(Chat.id == chat_id))
-        chat = result.scalars().first()
-
-        if not chat:
-            raise HTTPException(
-                status_code=404,
-                detail="Chat not found",
-            )
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            chat.search_space_id,
-            Permission.CHATS_READ.value,
-            "You don't have permission to read chats in this search space",
-        )
-
-        return chat
-    except HTTPException:
-        raise
-    except OperationalError:
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        raise HTTPException(
-            status_code=500,
-            detail="An unexpected error occurred while fetching the chat.",
-        ) from None
-
-
-@router.put("/chats/{chat_id}", response_model=ChatRead)
-async def update_chat(
-    chat_id: int,
-    chat_update: ChatUpdate,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Update a chat.
-    Requires CHATS_UPDATE permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Chat).filter(Chat.id == chat_id))
-        db_chat = result.scalars().first()
-
-        if not db_chat:
-            raise HTTPException(status_code=404, detail="Chat not found")
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            db_chat.search_space_id,
-            Permission.CHATS_UPDATE.value,
-            "You don't have permission to update chats in this search space",
-        )
-
-        update_data = chat_update.model_dump(exclude_unset=True)
-        for key, value in update_data.items():
-            if key == "messages":
-                db_chat.state_version = len(update_data["messages"])
-            setattr(db_chat, key, value)
-
-        await session.commit()
-        await session.refresh(db_chat)
-        return db_chat
-    except HTTPException:
-        raise
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400,
-            detail="Database constraint violation. Please check your input data.",
-        ) from None
-    except OperationalError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500,
-            detail="An unexpected error occurred while updating the chat.",
-        ) from None
-
-
-@router.delete("/chats/{chat_id}", response_model=dict)
-async def delete_chat(
-    chat_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Delete a chat.
-    Requires CHATS_DELETE permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Chat).filter(Chat.id == chat_id))
-        db_chat = result.scalars().first()
-
-        if not db_chat:
-            raise HTTPException(status_code=404, detail="Chat not found")
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            db_chat.search_space_id,
-            Permission.CHATS_DELETE.value,
-            "You don't have permission to delete chats in this search space",
-        )
-
-        await session.delete(db_chat)
-        await session.commit()
-        return {"message": "Chat deleted successfully"}
-    except HTTPException:
-        raise
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400, detail="Cannot delete chat due to existing dependencies."
-        ) from None
-    except OperationalError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=503, detail="Database operation failed. Please try again later."
-        ) from None
-    except Exception:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500,
-            detail="An unexpected error occurred while deleting the chat.",
-        ) from None

surfsense_backend/app/routes/new_chat_routes.py

@@ -13,6 +13,7 @@ These endpoints support the ThreadHistoryAdapter pattern from assistant-ui:
 from datetime import UTC, datetime
 
 from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import StreamingResponse
 from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
@@ -23,12 +24,14 @@ from app.db import (
     NewChatMessageRole,
     NewChatThread,
     Permission,
+    SearchSpace,
     User,
     get_async_session,
 )
 from app.schemas.new_chat import (
     NewChatMessageAppend,
     NewChatMessageRead,
+    NewChatRequest,
     NewChatThreadCreate,
     NewChatThreadRead,
     NewChatThreadUpdate,
@@ -37,6 +40,7 @@ from app.schemas.new_chat import (
     ThreadListItem,
     ThreadListResponse,
 )
+from app.tasks.chat.stream_new_chat import stream_new_chat
 from app.users import current_active_user
 from app.utils.rbac import check_permission
 
@@ -74,13 +78,10 @@ async def list_threads(
             "You don't have permission to read chats in this search space",
         )
 
-        # Get all threads for this user in this search space
+        # Get all threads in this search space
         query = (
             select(NewChatThread)
-            .filter(
+            .filter(NewChatThread.search_space_id == search_space_id)
-                NewChatThread.search_space_id == search_space_id,
-                NewChatThread.user_id == user.id,
-            )
             .order_by(NewChatThread.updated_at.desc())
         )
 
@@ -153,7 +154,6 @@ async def search_threads(
             select(NewChatThread)
             .filter(
                 NewChatThread.search_space_id == search_space_id,
-                NewChatThread.user_id == user.id,
                 NewChatThread.title.ilike(f"%{title}%"),
             )
             .order_by(NewChatThread.updated_at.desc())
@@ -211,7 +211,6 @@ async def create_thread(
             title=thread.title,
             archived=thread.archived,
             search_space_id=thread.search_space_id,
-            user_id=user.id,
             updated_at=now,
         )
         session.add(db_thread)
@@ -273,12 +272,6 @@ async def get_thread_messages(
             "You don't have permission to read chats in this search space",
         )
 
-        # Ensure user owns this thread
-        if thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         # Return messages in the format expected by assistant-ui
         messages = [
             NewChatMessageRead(
@@ -336,11 +329,6 @@ async def get_thread_full(
             "You don't have permission to read chats in this search space",
         )
 
-        if thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         return thread
 
     except HTTPException:
@@ -386,11 +374,6 @@ async def update_thread(
             "You don't have permission to update chats in this search space",
         )
 
-        if db_thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         # Update fields
         update_data = thread_update.model_dump(exclude_unset=True)
         for key, value in update_data.items():
@@ -451,11 +434,6 @@ async def delete_thread(
             "You don't have permission to delete chats in this search space",
         )
 
-        if db_thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         await session.delete(db_thread)
         await session.commit()
         return {"message": "Thread deleted successfully"}
@@ -530,11 +508,6 @@ async def append_message(
             "You don't have permission to update chats in this search space",
         )
 
-        if thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         # Convert string role to enum
         role_str = (
             message.role.lower() if isinstance(message.role, str) else message.role
@@ -639,11 +612,6 @@ async def list_messages(
             "You don't have permission to read chats in this search space",
         )
 
-        if thread.user_id != user.id:
-            raise HTTPException(
-                status_code=403, detail="You don't have access to this thread"
-            )
-
         # Get messages
         query = (
             select(NewChatMessage)
@@ -667,3 +635,79 @@ async def list_messages(
             status_code=500,
             detail=f"An unexpected error occurred while fetching messages: {e!s}",
         ) from None
+
+
+# =============================================================================
+# Chat Streaming Endpoint
+# =============================================================================
+
+
+@router.post("/new_chat")
+async def handle_new_chat(
+    request: NewChatRequest,
+    session: AsyncSession = Depends(get_async_session),
+    user: User = Depends(current_active_user),
+):
+    """
+    Stream chat responses from the deep agent.
+    
+    This endpoint handles the new chat functionality with streaming responses
+    using Server-Sent Events (SSE) format compatible with Vercel AI SDK.
+    
+    Requires CHATS_CREATE permission.
+    """
+    try:
+        # Verify thread exists and user has permission
+        result = await session.execute(
+            select(NewChatThread).filter(NewChatThread.id == request.chat_id)
+        )
+        thread = result.scalars().first()
+
+        if not thread:
+            raise HTTPException(status_code=404, detail="Thread not found")
+
+        await check_permission(
+            session,
+            user,
+            thread.search_space_id,
+            Permission.CHATS_CREATE.value,
+            "You don't have permission to chat in this search space",
+        )
+
+        # Get search space to check LLM config preferences
+        search_space_result = await session.execute(
+            select(SearchSpace).filter(SearchSpace.id == request.search_space_id)
+        )
+        search_space = search_space_result.scalars().first()
+
+        # Determine LLM config ID (use search space preference or default)
+        llm_config_id = -1  # Default to first global config
+        if search_space and search_space.fast_llm_id:
+            llm_config_id = search_space.fast_llm_id
+
+        # Return streaming response
+        return StreamingResponse(
+            stream_new_chat(
+                user_query=request.user_query,
+                user_id=str(user.id),
+                search_space_id=request.search_space_id,
+                chat_id=request.chat_id,
+                session=session,
+                llm_config_id=llm_config_id,
+                messages=request.messages,
+            ),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"An unexpected error occurred: {e!s}",
+        ) from None

surfsense_backend/app/routes/podcasts_routes.py

@@ -1,509 +0,0 @@
-import os
-from pathlib import Path
-
-from fastapi import APIRouter, Depends, HTTPException
-from fastapi.responses import StreamingResponse
-from sqlalchemy.exc import IntegrityError, SQLAlchemyError
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.future import select
-
-from app.db import (
-    Chat,
-    Permission,
-    Podcast,
-    SearchSpace,
-    SearchSpaceMembership,
-    User,
-    get_async_session,
-)
-from app.schemas import (
-    PodcastCreate,
-    PodcastGenerateRequest,
-    PodcastRead,
-    PodcastUpdate,
-)
-from app.tasks.podcast_tasks import generate_chat_podcast
-from app.users import current_active_user
-from app.utils.rbac import check_permission
-
-router = APIRouter()
-
-
-@router.post("/podcasts", response_model=PodcastRead)
-async def create_podcast(
-    podcast: PodcastCreate,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Create a new podcast.
-    Requires PODCASTS_CREATE permission.
-    """
-    try:
-        await check_permission(
-            session,
-            user,
-            podcast.search_space_id,
-            Permission.PODCASTS_CREATE.value,
-            "You don't have permission to create podcasts in this search space",
-        )
-        db_podcast = Podcast(**podcast.model_dump())
-        session.add(db_podcast)
-        await session.commit()
-        await session.refresh(db_podcast)
-        return db_podcast
-    except HTTPException as he:
-        raise he
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400,
-            detail="Podcast creation failed due to constraint violation",
-        ) from None
-    except SQLAlchemyError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while creating podcast"
-        ) from None
-    except Exception:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail="An unexpected error occurred"
-        ) from None
-
-
-@router.get("/podcasts", response_model=list[PodcastRead])
-async def read_podcasts(
-    skip: int = 0,
-    limit: int = 100,
-    search_space_id: int | None = None,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    List podcasts the user has access to.
-    Requires PODCASTS_READ permission for the search space(s).
-    """
-    if skip < 0 or limit < 1:
-        raise HTTPException(status_code=400, detail="Invalid pagination parameters")
-    try:
-        if search_space_id is not None:
-            # Check permission for specific search space
-            await check_permission(
-                session,
-                user,
-                search_space_id,
-                Permission.PODCASTS_READ.value,
-                "You don't have permission to read podcasts in this search space",
-            )
-            result = await session.execute(
-                select(Podcast)
-                .filter(Podcast.search_space_id == search_space_id)
-                .offset(skip)
-                .limit(limit)
-            )
-        else:
-            # Get podcasts from all search spaces user has membership in
-            result = await session.execute(
-                select(Podcast)
-                .join(SearchSpace)
-                .join(SearchSpaceMembership)
-                .filter(SearchSpaceMembership.user_id == user.id)
-                .offset(skip)
-                .limit(limit)
-            )
-        return result.scalars().all()
-    except HTTPException:
-        raise
-    except SQLAlchemyError:
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while fetching podcasts"
-        ) from None
-
-
-@router.get("/podcasts/{podcast_id}", response_model=PodcastRead)
-async def read_podcast(
-    podcast_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Get a specific podcast by ID.
-    Requires PODCASTS_READ permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Podcast).filter(Podcast.id == podcast_id))
-        podcast = result.scalars().first()
-
-        if not podcast:
-            raise HTTPException(
-                status_code=404,
-                detail="Podcast not found",
-            )
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            podcast.search_space_id,
-            Permission.PODCASTS_READ.value,
-            "You don't have permission to read podcasts in this search space",
-        )
-
-        return podcast
-    except HTTPException as he:
-        raise he
-    except SQLAlchemyError:
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while fetching podcast"
-        ) from None
-
-
-@router.put("/podcasts/{podcast_id}", response_model=PodcastRead)
-async def update_podcast(
-    podcast_id: int,
-    podcast_update: PodcastUpdate,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Update a podcast.
-    Requires PODCASTS_UPDATE permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Podcast).filter(Podcast.id == podcast_id))
-        db_podcast = result.scalars().first()
-
-        if not db_podcast:
-            raise HTTPException(status_code=404, detail="Podcast not found")
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            db_podcast.search_space_id,
-            Permission.PODCASTS_UPDATE.value,
-            "You don't have permission to update podcasts in this search space",
-        )
-
-        update_data = podcast_update.model_dump(exclude_unset=True)
-        for key, value in update_data.items():
-            setattr(db_podcast, key, value)
-        await session.commit()
-        await session.refresh(db_podcast)
-        return db_podcast
-    except HTTPException as he:
-        raise he
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400, detail="Update failed due to constraint violation"
-        ) from None
-    except SQLAlchemyError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while updating podcast"
-        ) from None
-
-
-@router.delete("/podcasts/{podcast_id}", response_model=dict)
-async def delete_podcast(
-    podcast_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Delete a podcast.
-    Requires PODCASTS_DELETE permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Podcast).filter(Podcast.id == podcast_id))
-        db_podcast = result.scalars().first()
-
-        if not db_podcast:
-            raise HTTPException(status_code=404, detail="Podcast not found")
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            db_podcast.search_space_id,
-            Permission.PODCASTS_DELETE.value,
-            "You don't have permission to delete podcasts in this search space",
-        )
-
-        await session.delete(db_podcast)
-        await session.commit()
-        return {"message": "Podcast deleted successfully"}
-    except HTTPException as he:
-        raise he
-    except SQLAlchemyError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while deleting podcast"
-        ) from None
-
-
-async def generate_chat_podcast_with_new_session(
-    chat_id: int,
-    search_space_id: int,
-    user_id: int,
-    podcast_title: str | None = None,
-    user_prompt: str | None = None,
-):
-    """Create a new session and process chat podcast generation."""
-    from app.db import async_session_maker
-
-    async with async_session_maker() as session:
-        try:
-            await generate_chat_podcast(
-                session, chat_id, search_space_id, user_id, podcast_title, user_prompt
-            )
-        except Exception as e:
-            import logging
-
-            logging.error(f"Error generating podcast from chat: {e!s}")
-
-
-@router.post("/podcasts/generate")
-async def generate_podcast(
-    request: PodcastGenerateRequest,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Generate a podcast from a chat or document.
-    Requires PODCASTS_CREATE permission.
-    """
-    try:
-        # Check if the user has permission to create podcasts
-        await check_permission(
-            session,
-            user,
-            request.search_space_id,
-            Permission.PODCASTS_CREATE.value,
-            "You don't have permission to create podcasts in this search space",
-        )
-
-        if request.type == "CHAT":
-            # Verify that all chat IDs belong to this user and search space
-            query = (
-                select(Chat)
-                .filter(
-                    Chat.id.in_(request.ids),
-                    Chat.search_space_id == request.search_space_id,
-                )
-                .join(SearchSpace)
-                .filter(SearchSpace.user_id == user.id)
-            )
-
-            result = await session.execute(query)
-            valid_chats = result.scalars().all()
-            valid_chat_ids = [chat.id for chat in valid_chats]
-
-            # If any requested ID is not in valid IDs, raise error immediately
-            if len(valid_chat_ids) != len(request.ids):
-                raise HTTPException(
-                    status_code=403,
-                    detail="One or more chat IDs do not belong to this user or search space",
-                )
-
-            from app.tasks.celery_tasks.podcast_tasks import (
-                generate_chat_podcast_task,
-            )
-
-            # Add Celery tasks for each chat ID
-            for chat_id in valid_chat_ids:
-                generate_chat_podcast_task.delay(
-                    chat_id,
-                    request.search_space_id,
-                    user.id,
-                    request.podcast_title,
-                    request.user_prompt,
-                )
-
-        return {
-            "message": "Podcast generation started",
-        }
-    except HTTPException as he:
-        raise he
-    except IntegrityError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=400,
-            detail="Podcast generation failed due to constraint violation",
-        ) from None
-    except SQLAlchemyError:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail="Database error occurred while generating podcast"
-        ) from None
-    except Exception as e:
-        await session.rollback()
-        raise HTTPException(
-            status_code=500, detail=f"An unexpected error occurred: {e!s}"
-        ) from e
-
-
-@router.get("/podcasts/{podcast_id}/stream")
-async def stream_podcast(
-    podcast_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Stream a podcast audio file.
-    Requires PODCASTS_READ permission for the search space.
-    """
-    try:
-        result = await session.execute(select(Podcast).filter(Podcast.id == podcast_id))
-        podcast = result.scalars().first()
-
-        if not podcast:
-            raise HTTPException(
-                status_code=404,
-                detail="Podcast not found",
-            )
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            podcast.search_space_id,
-            Permission.PODCASTS_READ.value,
-            "You don't have permission to access podcasts in this search space",
-        )
-
-        # Get the file path
-        file_path = podcast.file_location
-
-        # Check if the file exists
-        if not os.path.isfile(file_path):
-            raise HTTPException(status_code=404, detail="Podcast audio file not found")
-
-        # Define a generator function to stream the file
-        def iterfile():
-            with open(file_path, mode="rb") as file_like:
-                yield from file_like
-
-        # Return a streaming response with appropriate headers
-        return StreamingResponse(
-            iterfile(),
-            media_type="audio/mpeg",
-            headers={
-                "Accept-Ranges": "bytes",
-                "Content-Disposition": f"inline; filename={Path(file_path).name}",
-            },
-        )
-
-    except HTTPException as he:
-        raise he
-    except Exception as e:
-        raise HTTPException(
-            status_code=500, detail=f"Error streaming podcast: {e!s}"
-        ) from e
-
-
-@router.get("/podcasts/by-chat/{chat_id}", response_model=PodcastRead | None)
-async def get_podcast_by_chat_id(
-    chat_id: int,
-    session: AsyncSession = Depends(get_async_session),
-    user: User = Depends(current_active_user),
-):
-    """
-    Get a podcast by its associated chat ID.
-    Requires PODCASTS_READ permission for the search space.
-    """
-    try:
-        # First get the chat to find its search space
-        chat_result = await session.execute(select(Chat).filter(Chat.id == chat_id))
-        chat = chat_result.scalars().first()
-
-        if not chat:
-            return None
-
-        # Check permission for the search space
-        await check_permission(
-            session,
-            user,
-            chat.search_space_id,
-            Permission.PODCASTS_READ.value,
-            "You don't have permission to read podcasts in this search space",
-        )
-
-        # Get the podcast
-        result = await session.execute(
-            select(Podcast).filter(Podcast.chat_id == chat_id)
-        )
-        podcast = result.scalars().first()
-
-        return podcast
-    except HTTPException as he:
-        raise he
-    except Exception as e:
-        raise HTTPException(
-            status_code=500, detail=f"Error fetching podcast: {e!s}"
-        ) from e
-
-
-@router.get("/podcasts/task/{task_id}/status")
-async def get_podcast_task_status(
-    task_id: str,
-    user: User = Depends(current_active_user),
-):
-    """
-    Get the status of a podcast generation task.
-    Used by new-chat frontend to poll for completion.
-
-    Returns:
-    - status: "processing" | "success" | "error"
-    - podcast_id: (only if status == "success")
-    - title: (only if status == "success")
-    - error: (only if status == "error")
-    """
-    try:
-        from celery.result import AsyncResult
-
-        from app.celery_app import celery_app
-
-        result = AsyncResult(task_id, app=celery_app)
-
-        if result.ready():
-            # Task completed
-            if result.successful():
-                task_result = result.result
-                if isinstance(task_result, dict):
-                    if task_result.get("status") == "success":
-                        return {
-                            "status": "success",
-                            "podcast_id": task_result.get("podcast_id"),
-                            "title": task_result.get("title"),
-                            "transcript_entries": task_result.get("transcript_entries"),
-                        }
-                    else:
-                        return {
-                            "status": "error",
-                            "error": task_result.get("error", "Unknown error"),
-                        }
-                else:
-                    return {
-                        "status": "error",
-                        "error": "Unexpected task result format",
-                    }
-            else:
-                # Task failed
-                return {
-                    "status": "error",
-                    "error": str(result.result) if result.result else "Task failed",
-                }
-        else:
-            # Task still processing
-            return {
-                "status": "processing",
-                "state": result.state,
-            }
-
-    except Exception as e:
-        raise HTTPException(
-            status_code=500, detail=f"Error checking task status: {e!s}"
-        ) from e

surfsense_backend/app/schemas/__init__.py

@@ -1,13 +1,4 @@
 from .base import IDModel, TimestampModel
-from .chats import (
-    AISDKChatRequest,
-    ChatBase,
-    ChatCreate,
-    ChatRead,
-    ChatReadWithoutMessages,
-    ChatUpdate,
-    NewChatRequest,
-)
 from .chunks import ChunkBase, ChunkCreate, ChunkRead, ChunkUpdate
 from .documents import (
     DocumentBase,
@@ -22,9 +13,11 @@ from .documents import (
 from .llm_config import LLMConfigBase, LLMConfigCreate, LLMConfigRead, LLMConfigUpdate
 from .logs import LogBase, LogCreate, LogFilter, LogRead, LogUpdate
 from .new_chat import (
+    ChatMessage,
     NewChatMessageAppend,
     NewChatMessageCreate,
     NewChatMessageRead,
+    NewChatRequest,
     NewChatThreadCreate,
     NewChatThreadRead,
     NewChatThreadUpdate,
@@ -33,13 +26,6 @@ from .new_chat import (
     ThreadListItem,
     ThreadListResponse,
 )
-from .podcasts import (
-    PodcastBase,
-    PodcastCreate,
-    PodcastGenerateRequest,
-    PodcastRead,
-    PodcastUpdate,
-)
 from .rbac_schemas import (
     InviteAcceptRequest,
     InviteAcceptResponse,
@@ -73,44 +59,8 @@ from .search_space import (
 from .users import UserCreate, UserRead, UserUpdate
 
 __all__ = [
-    "AISDKChatRequest",
+    # Chat schemas (assistant-ui integration)
-    "ChatBase",
+    "ChatMessage",
-    "ChatCreate",
-    "ChatRead",
-    "ChatReadWithoutMessages",
-    "ChatUpdate",
-    "ChunkBase",
-    "ChunkCreate",
-    "ChunkRead",
-    "ChunkUpdate",
-    "DocumentBase",
-    "DocumentRead",
-    "DocumentUpdate",
-    "DocumentWithChunksRead",
-    "DocumentsCreate",
-    "ExtensionDocumentContent",
-    "ExtensionDocumentMetadata",
-    "IDModel",
-    # RBAC schemas
-    "InviteAcceptRequest",
-    "InviteAcceptResponse",
-    "InviteCreate",
-    "InviteInfoResponse",
-    "InviteRead",
-    "InviteUpdate",
-    "LLMConfigBase",
-    "LLMConfigCreate",
-    "LLMConfigRead",
-    "LLMConfigUpdate",
-    "LogBase",
-    "LogCreate",
-    "LogFilter",
-    "LogRead",
-    "LogUpdate",
-    "MembershipRead",
-    "MembershipReadWithUser",
-    "MembershipUpdate",
-    # New chat schemas (assistant-ui integration)
     "NewChatMessageAppend",
     "NewChatMessageCreate",
     "NewChatMessageRead",
@@ -119,30 +69,64 @@ __all__ = [
     "NewChatThreadRead",
     "NewChatThreadUpdate",
     "NewChatThreadWithMessages",
+    "ThreadHistoryLoadResponse",
+    "ThreadListItem",
+    "ThreadListResponse",
+    # Chunk schemas
+    "ChunkBase",
+    "ChunkCreate",
+    "ChunkRead",
+    "ChunkUpdate",
+    # Document schemas
+    "DocumentBase",
+    "DocumentRead",
+    "DocumentUpdate",
+    "DocumentWithChunksRead",
+    "DocumentsCreate",
+    "ExtensionDocumentContent",
+    "ExtensionDocumentMetadata",
     "PaginatedResponse",
+    # Base schemas
+    "IDModel",
+    "TimestampModel",
+    # LLM Config schemas
+    "LLMConfigBase",
+    "LLMConfigCreate",
+    "LLMConfigRead",
+    "LLMConfigUpdate",
+    # Log schemas
+    "LogBase",
+    "LogCreate",
+    "LogFilter",
+    "LogRead",
+    "LogUpdate",
+    # RBAC schemas
+    "InviteAcceptRequest",
+    "InviteAcceptResponse",
+    "InviteCreate",
+    "InviteInfoResponse",
+    "InviteRead",
+    "InviteUpdate",
+    "MembershipRead",
+    "MembershipReadWithUser",
+    "MembershipUpdate",
     "PermissionInfo",
     "PermissionsListResponse",
-    "PodcastBase",
-    "PodcastCreate",
-    "PodcastGenerateRequest",
-    "PodcastRead",
-    "PodcastUpdate",
     "RoleCreate",
     "RoleRead",
     "RoleUpdate",
+    # Search source connector schemas
     "SearchSourceConnectorBase",
     "SearchSourceConnectorCreate",
     "SearchSourceConnectorRead",
     "SearchSourceConnectorUpdate",
+    # Search space schemas
     "SearchSpaceBase",
     "SearchSpaceCreate",
     "SearchSpaceRead",
     "SearchSpaceUpdate",
     "SearchSpaceWithStats",
-    "ThreadHistoryLoadResponse",
+    # User schemas
-    "ThreadListItem",
-    "ThreadListResponse",
-    "TimestampModel",
     "UserCreate",
     "UserRead",
     "UserSearchSpaceAccess",

surfsense_backend/app/schemas/chats.py

@@ -1,80 +0,0 @@
-from typing import Any
-
-from pydantic import BaseModel, ConfigDict
-
-from app.db import ChatType
-
-from .base import IDModel, TimestampModel
-
-
-class ChatBase(BaseModel):
-    type: ChatType
-    title: str
-    initial_connectors: list[str] | None = None
-    messages: list[Any]
-    search_space_id: int
-    state_version: int = 1
-
-
-class ChatBaseWithoutMessages(BaseModel):
-    type: ChatType
-    title: str
-    search_space_id: int
-    state_version: int = 1
-
-
-class ClientAttachment(BaseModel):
-    name: str
-    content_type: str
-    url: str
-
-
-class ToolInvocation(BaseModel):
-    tool_call_id: str
-    tool_name: str
-    args: dict
-    result: dict
-
-
-# class ClientMessage(BaseModel):
-#     role: str
-#     content: str
-#     experimental_attachments: Optional[List[ClientAttachment]] = None
-#     toolInvocations: Optional[List[ToolInvocation]] = None
-
-
-class AISDKChatRequest(BaseModel):
-    messages: list[Any]
-    data: dict[str, Any] | None = None
-
-
-class ChatMessage(BaseModel):
-    """A single message in the chat history."""
-
-    role: str  # "user" or "assistant"
-    content: str
-
-
-class NewChatRequest(BaseModel):
-    """Request schema for the new deep agent chat endpoint."""
-
-    chat_id: int
-    user_query: str
-    search_space_id: int
-    messages: list[ChatMessage] | None = None  # Optional chat history from frontend
-
-
-class ChatCreate(ChatBase):
-    pass
-
-
-class ChatUpdate(ChatBase):
-    pass
-
-
-class ChatRead(ChatBase, IDModel, TimestampModel):
-    model_config = ConfigDict(from_attributes=True)
-
-
-class ChatReadWithoutMessages(ChatBaseWithoutMessages, IDModel, TimestampModel):
-    model_config = ConfigDict(from_attributes=True)

surfsense_backend/app/schemas/new_chat.py

@@ -127,3 +127,24 @@ class ThreadListResponse(BaseModel):
 
     threads: list[ThreadListItem]
     archived_threads: list[ThreadListItem]
+
+
+# =============================================================================
+# Chat Request Schemas (for deep agent)
+# =============================================================================
+
+
+class ChatMessage(BaseModel):
+    """A single message in the chat history."""
+
+    role: str  # "user" or "assistant"
+    content: str
+
+
+class NewChatRequest(BaseModel):
+    """Request schema for the deep agent chat endpoint."""
+
+    chat_id: int
+    user_query: str
+    search_space_id: int
+    messages: list[ChatMessage] | None = None  # Optional chat history from frontend

surfsense_backend/app/schemas/podcasts.py

@@ -1,33 +0,0 @@
-from typing import Any, Literal
-
-from pydantic import BaseModel, ConfigDict
-
-from .base import IDModel, TimestampModel
-
-
-class PodcastBase(BaseModel):
-    title: str
-    podcast_transcript: list[Any]
-    file_location: str = ""
-    search_space_id: int
-    chat_state_version: int | None = None
-
-
-class PodcastCreate(PodcastBase):
-    pass
-
-
-class PodcastUpdate(PodcastBase):
-    pass
-
-
-class PodcastRead(PodcastBase, IDModel, TimestampModel):
-    model_config = ConfigDict(from_attributes=True)
-
-
-class PodcastGenerateRequest(BaseModel):
-    type: Literal["DOCUMENT", "CHAT"]
-    ids: list[int]
-    search_space_id: int
-    podcast_title: str | None = None
-    user_prompt: str | None = None

surfsense_backend/app/tasks/chat/stream_connector_search_results.py

@@ -1,75 +0,0 @@
-from collections.abc import AsyncGenerator
-from typing import Any
-from uuid import UUID
-
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from app.agents.researcher.graph import graph as researcher_graph
-from app.agents.researcher.state import State
-from app.services.streaming_service import StreamingService
-
-
-async def stream_connector_search_results(
-    user_query: str,
-    user_id: str | UUID,
-    search_space_id: int,
-    session: AsyncSession,
-    research_mode: str,
-    selected_connectors: list[str],
-    langchain_chat_history: list[Any],
-    document_ids_to_add_in_context: list[int],
-    language: str | None = None,
-    top_k: int = 10,
-) -> AsyncGenerator[str, None]:
-    """
-    Stream connector search results to the client
-
-    Args:
-        user_query: The user's query
-        user_id: The user's ID (can be UUID object or string)
-        search_space_id: The search space ID
-        session: The database session
-        research_mode: The research mode
-        selected_connectors: List of selected connectors
-
-    Yields:
-        str: Formatted response strings
-    """
-    streaming_service = StreamingService()
-
-    # Convert UUID to string if needed
-    user_id_str = str(user_id) if isinstance(user_id, UUID) else user_id
-
-    # Sample configuration
-    config = {
-        "configurable": {
-            "user_query": user_query,
-            "connectors_to_search": selected_connectors,
-            "user_id": user_id_str,
-            "search_space_id": search_space_id,
-            "document_ids_to_add_in_context": document_ids_to_add_in_context,
-            "language": language,  # Add language to the configuration
-            "top_k": top_k,  # Add top_k to the configuration
-        }
-    }
-    # print(f"Researcher configuration: {config['configurable']}")  # Debug print
-    # Initialize state with database session and streaming service
-    initial_state = State(
-        db_session=session,
-        streaming_service=streaming_service,
-        chat_history=langchain_chat_history,
-    )
-
-    # Run the graph directly
-    print("\nRunning the complete researcher workflow...")
-
-    # Use streaming with config parameter
-    async for chunk in researcher_graph.astream(
-        initial_state,
-        config=config,
-        stream_mode="custom",
-    ):
-        if isinstance(chunk, dict) and "yield_value" in chunk:
-            yield chunk["yield_value"]
-
-    yield streaming_service.format_completion()

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -18,7 +18,7 @@ from app.agents.new_chat.llm_config import (
     create_chat_litellm_from_config,
     load_llm_config_from_yaml,
 )
-from app.schemas.chats import ChatMessage
+from app.schemas.new_chat import ChatMessage
 from app.services.connector_service import ConnectorService
 from app.services.new_streaming_service import VercelStreamingService