Add summary column and enhance search vector generation for embeddings

2025-12-25 07:59:35 +01:00 · 2025-02-22 18:43:56 +01:00
parent 29c96f7dd6
commit c33fc4a024
15 changed files with 356 additions and 131 deletions
--- a/apps/server/src/data/migrations/00018-create-node-embeddings-table.ts
+++ b/apps/server/src/data/migrations/00018-create-node-embeddings-table.ts
@@ -10,10 +10,14 @@ export const createNodeEmbeddingsTable: Migration = {
      .addColumn('root_id', 'varchar(30)', (col) => col.notNull())
      .addColumn('workspace_id', 'varchar(30)', (col) => col.notNull())
      .addColumn('text', 'text', (col) => col.notNull())
+      .addColumn('summary', 'text')
      .addColumn('embedding_vector', sql`vector(2000)`, (col) => col.notNull())
      .addColumn(
        'search_vector',
-        sql`tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED`
+        sql`tsvector GENERATED ALWAYS AS (
+          setweight(to_tsvector('english', COALESCE(text, '')), 'A') ||
+          setweight(to_tsvector('english', COALESCE(summary, '')), 'B')
+        ) STORED`
      )
      .addColumn('created_at', 'timestamptz', (col) => col.notNull())
      .addColumn('updated_at', 'timestamptz')
--- a/apps/server/src/data/migrations/00019-create-document-embeddings-table.ts
+++ b/apps/server/src/data/migrations/00019-create-document-embeddings-table.ts
@@ -8,10 +8,14 @@ export const createDocumentEmbeddingsTable: Migration = {
      .addColumn('chunk', 'integer', (col) => col.notNull())
      .addColumn('workspace_id', 'varchar(30)', (col) => col.notNull())
      .addColumn('text', 'text', (col) => col.notNull())
+      .addColumn('summary', 'text')
      .addColumn('embedding_vector', sql`vector(2000)`, (col) => col.notNull())
      .addColumn(
        'search_vector',
-        sql`tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED`
+        sql`tsvector GENERATED ALWAYS AS (
+          setweight(to_tsvector('english', COALESCE(text, '')), 'A') ||
+          setweight(to_tsvector('english', COALESCE(summary, '')), 'B')
+        ) STORED`
      )
      .addColumn('created_at', 'timestamptz', (col) => col.notNull())
      .addColumn('updated_at', 'timestamptz')
--- a/apps/server/src/data/schema.ts
+++ b/apps/server/src/data/schema.ts
@@ -265,6 +265,7 @@ interface NodeEmbeddingTable {
  root_id: ColumnType<string, string, never>;
  workspace_id: ColumnType<string, string, never>;
  text: ColumnType<string, string, string>;
+  summary: ColumnType<string | null, string | null, string | null>;
  embedding_vector: ColumnType<number[], number[], number[]>;
  search_vector: ColumnType<never, never, never>;
  created_at: ColumnType<Date, Date, never>;
@@ -280,6 +281,7 @@ interface DocumentEmbeddingTable {
  chunk: ColumnType<number, number, number>;
  workspace_id: ColumnType<string, string, never>;
  text: ColumnType<string, string, string>;
+  summary: ColumnType<string | null, string | null, string | null>;
  embedding_vector: ColumnType<number[], number[], number[]>;
  search_vector: ColumnType<never, never, never>;
  created_at: ColumnType<Date, Date, never>;
--- a/apps/server/src/jobs/assistant-response.ts
+++ b/apps/server/src/jobs/assistant-response.ts
@@ -81,7 +81,7 @@ export const assistantResponseHandler: JobHandler<
      userId: user.id,
      userDetails: {
        name: user.name || 'User',
-        email: user.email || 'unknown@example.com',
+        email: user.email || '',
      },
      parentMessageId: message.parent_id || message.id,
      currentMessageId: message.id,
--- a/apps/server/src/jobs/embed-document.ts
+++ b/apps/server/src/jobs/embed-document.ts
@@ -59,14 +59,6 @@ export const embedDocumentHandler = async (input: {
    return;
  }

-  const textChunks = await chunkText(
-    text,
-    {
-      type: 'document',
-      node: node,
-    },
-    enrichChunk
-  );
  const embeddings = new OpenAIEmbeddings({
    apiKey: configuration.ai.embedding.apiKey,
    modelName: configuration.ai.embedding.modelName,
@@ -75,21 +67,38 @@ export const embedDocumentHandler = async (input: {

  const existingEmbeddings = await database
    .selectFrom('document_embeddings')
-    .select(['chunk', 'text'])
+    .select(['chunk', 'text', 'summary'])
    .where('document_id', '=', documentId)
    .execute();

+  const textChunks = await chunkText(
+    text,
+    existingEmbeddings.map((e) => ({
+      text: e.text,
+      summary: e.summary ?? undefined,
+    })),
+    { type: 'document', node: node },
+    enrichChunk
+  );
+
  const embeddingsToCreateOrUpdate: CreateDocumentEmbedding[] = [];
  for (let i = 0; i < textChunks.length; i++) {
-    const textChunk = textChunks[i];
-    if (!textChunk) continue;
+    const chunk = textChunks[i];
+    if (!chunk) {
+      continue;
+    }
+
    const existing = existingEmbeddings.find((e) => e.chunk === i);
-    if (existing && existing.text === textChunk) continue;
+    if (existing && existing.text === chunk.text) {
+      continue;
+    }
+
    embeddingsToCreateOrUpdate.push({
      document_id: documentId,
      chunk: i,
      workspace_id: document.workspace_id,
-      text: textChunk,
+      text: chunk.text,
+      summary: chunk.summary,
      embedding_vector: [],
      created_at: new Date(),
    });
@@ -98,7 +107,9 @@ export const embedDocumentHandler = async (input: {
  const batchSize = configuration.ai.embedding.batchSize;
  for (let i = 0; i < embeddingsToCreateOrUpdate.length; i += batchSize) {
    const batch = embeddingsToCreateOrUpdate.slice(i, i + batchSize);
-    const textsToEmbed = batch.map((item) => item.text);
+    const textsToEmbed = batch.map((item) =>
+      item.summary ? `${item.summary}\n\n${item.text}` : item.text
+    );
    const embeddingVectors = await embeddings.embedDocuments(textsToEmbed);
    for (let j = 0; j < batch.length; j++) {
      const vector = embeddingVectors[j];
@@ -121,6 +132,7 @@ export const embedDocumentHandler = async (input: {
        chunk: embedding.chunk,
        workspace_id: embedding.workspace_id,
        text: embedding.text,
+        summary: embedding.summary,
        embedding_vector: sql.raw(
          `'[${embedding.embedding_vector.join(',')}]'::vector`
        ),
@@ -129,6 +141,7 @@ export const embedDocumentHandler = async (input: {
      .onConflict((oc) =>
        oc.columns(['document_id', 'chunk']).doUpdateSet({
          text: sql.ref('excluded.text'),
+          summary: sql.ref('excluded.summary'),
          embedding_vector: sql.ref('excluded.embedding_vector'),
          updated_at: new Date(),
        })
--- a/apps/server/src/jobs/embed-node.ts
+++ b/apps/server/src/jobs/embed-node.ts
@@ -154,12 +154,6 @@ export const embedNodeHandler = async (input: {
    return;
  }

-  const textChunks = await chunkText(
-    text,
-    { type: 'node', node: node },
-    enrichChunk
-  );
-
  const embeddings = new OpenAIEmbeddings({
    apiKey: configuration.ai.embedding.apiKey,
    modelName: configuration.ai.embedding.modelName,
@@ -168,19 +162,29 @@ export const embedNodeHandler = async (input: {

  const existingEmbeddings = await database
    .selectFrom('node_embeddings')
-    .select(['chunk', 'text'])
+    .select(['chunk', 'text', 'summary'])
    .where('node_id', '=', nodeId)
    .execute();

+  const textChunks = await chunkText(
+    text,
+    existingEmbeddings.map((e) => ({
+      text: e.text,
+      summary: e.summary ?? undefined,
+    })),
+    { type: 'node', node: node },
+    enrichChunk
+  );
+
  const embeddingsToCreateOrUpdate: CreateNodeEmbedding[] = [];
  for (let i = 0; i < textChunks.length; i++) {
-    const textChunk = textChunks[i];
-    if (!textChunk) {
+    const chunk = textChunks[i];
+    if (!chunk) {
      continue;
    }

    const existing = existingEmbeddings.find((e) => e.chunk === i);
-    if (existing && existing.text === textChunk) {
+    if (existing && existing.text === chunk.text) {
      continue;
    }

@@ -190,7 +194,8 @@ export const embedNodeHandler = async (input: {
      parent_id: node.parent_id,
      root_id: node.root_id,
      workspace_id: node.workspace_id,
-      text: textChunk,
+      text: chunk.text,
+      summary: chunk.summary,
      embedding_vector: [],
      created_at: new Date(),
    });
@@ -199,7 +204,9 @@ export const embedNodeHandler = async (input: {
  const batchSize = configuration.ai.embedding.batchSize;
  for (let i = 0; i < embeddingsToCreateOrUpdate.length; i += batchSize) {
    const batch = embeddingsToCreateOrUpdate.slice(i, i + batchSize);
-    const textsToEmbed = batch.map((item) => item.text);
+    const textsToEmbed = batch.map((item) =>
+      item.summary ? `${item.summary}\n\n${item.text}` : item.text
+    );
    const embeddingVectors = await embeddings.embedDocuments(textsToEmbed);
    for (let j = 0; j < batch.length; j++) {
      const vector = embeddingVectors[j];
@@ -224,6 +231,7 @@ export const embedNodeHandler = async (input: {
        root_id: embedding.root_id,
        workspace_id: embedding.workspace_id,
        text: embedding.text,
+        summary: embedding.summary,
        embedding_vector: sql.raw(
          `'[${embedding.embedding_vector.join(',')}]'::vector`
        ),
@@ -232,6 +240,7 @@ export const embedNodeHandler = async (input: {
      .onConflict((oc) =>
        oc.columns(['node_id', 'chunk']).doUpdateSet({
          text: sql.ref('excluded.text'),
+          summary: sql.ref('excluded.summary'),
          embedding_vector: sql.ref('excluded.embedding_vector'),
          updated_at: new Date(),
        })
--- a/apps/server/src/lib/assistant.ts
+++ b/apps/server/src/lib/assistant.ts
@@ -166,7 +166,7 @@ async function rerankContextDocuments(state: AssistantChainState) {
  }));
  const rerankedContext = await rerankDocuments(
    docsForRerank,
-    state.rewrittenQuery
+    state.rewrittenQuery.semanticQuery
  );

  return { rerankedContext };
--- a/apps/server/src/lib/chunking.ts
+++ b/apps/server/src/lib/chunking.ts
@@ -224,12 +224,13 @@ async function fetchMetadata(metadata?: {

 export async function chunkText(
  text: string,
+  existingChunks: Array<{ text: string; summary?: string }>,
  metadata?: { type: 'node' | 'document'; node: SelectNode },
  enrichFn?: (
    prompt: string,
    baseVars: Record<string, string>
  ) => Promise<string>
-): Promise<string[]> {
+): Promise<Array<{ text: string; summary?: string }>> {
  const chunkSize = configuration.ai.chunking.defaultChunkSize;
  const chunkOverlap = configuration.ai.chunking.defaultOverlap;
  const splitter = new RecursiveCharacterTextSplitter({
@@ -238,26 +239,38 @@ export async function chunkText(
  });
  const docs = await splitter.createDocuments([text]);
  let chunks = docs
-    .map((doc) => doc.pageContent)
-    .filter((c) => c.trim().length > 5);
+    .map((doc) => ({ text: doc.pageContent }))
+    .filter((c) => c.text.trim().length > 5);

  if (configuration.ai.chunking.enhanceWithContext && enrichFn && metadata) {
-    const enriched: string[] = [];
    const enrichedMetadata = await fetchMetadata(metadata);
    if (!enrichedMetadata) {
      return chunks;
    }
+
+    const enrichedChunks: Array<{ text: string; summary: string }> = [];
+
    for (const chunk of chunks) {
+      const existingChunk = existingChunks.find((ec) => ec.text === chunk.text);
+      if (existingChunk?.summary) {
+        enrichedChunks.push({
+          text: chunk.text,
+          summary: existingChunk.summary,
+        });
+
+        continue;
+      }
+
      const { prompt, baseVars } = prepareEnrichmentPrompt(
-        chunk,
+        chunk.text,
        text,
        enrichedMetadata
      );

-      const enrichment = await enrichFn(prompt, baseVars);
-      enriched.push(enrichment + '\n\n' + chunk);
+      const summary = await enrichFn(prompt, baseVars);
+      enrichedChunks.push({ text: chunk.text, summary });
    }
-    return enriched;
+    return enrichedChunks;
  }
  return chunks;
 }
--- a/apps/server/src/lib/llm-prompts.ts
+++ b/apps/server/src/lib/llm-prompts.ts
@@ -2,94 +2,244 @@ import { PromptTemplate, ChatPromptTemplate } from '@langchain/core/prompts';
 import { NodeMetadata, DocumentMetadata } from '@/types/chunking';

 export const queryRewritePrompt = PromptTemplate.fromTemplate(
-  `You are an expert at rewriting queries for information retrieval within Colanode.
+  `<task>
+  You are an expert at rewriting search queries to optimize for both semantic similarity and keyword-based search in a document retrieval system.
+  Your task is to generate two separate optimized queries:
+  1. A semantic search query optimized for vector embeddings and semantic similarity
+  2. A keyword search query optimized for full-text search using PostgreSQL's tsquery
+</task>
+
+<guidelines>
+  For semantic search query:
+  1. Focus on conceptual meaning and intent
+  2. Include context-indicating terms
+  3. Preserve relationship words between concepts
+  4. Expand concepts with related terms
+  5. Remove noise words and syntax-specific terms
  
-Guidelines:
-1. Extract the core information need.
-2. Remove filler words.
-3. Preserve key technical terms and dates.
-  
-Original query:
-{query}
-  
-Rewrite the query and return only the rewritten version.`
+  For keyword search query:
+  1. Focus on specific technical terms and exact matches
+  2. Include variations of key terms
+  3. Keep proper nouns and domain-specific vocabulary
+  4. Optimize for PostgreSQL's websearch_to_tsquery syntax
+  5. Include essential filters and constraints
+</guidelines>
+
+<input>
+  Original query: {query}
+</input>
+
+<output_format>
+  Return a JSON object with:
+  {{
+    "semanticQuery": "optimized query for semantic search",
+    "keywordQuery": "optimized query for keyword search"
+}}
+</output_format>`
 );

 export const summarizationPrompt = PromptTemplate.fromTemplate(
-  `Summarize the following text focusing on key points relevant to the user's query.
-If the text is short (<100 characters), return it as is.
-  
-Text: {text}
-User Query: {query}`
+  `<task>
+  Summarize the following text focusing on key points relevant to the user's query.
+  If the text is short (<100 characters), return it as is.
+</task>
+
+<input>
+  Text: {text}
+  User Query: {query}
+</input>`
 );

 export const rerankPrompt = PromptTemplate.fromTemplate(
-  `Re-rank the following list of documents by their relevance to the query.
-For each document, provide:
- Original index (from input)
- A relevance score between 0 and 1
- Document type (node or document)
- Source ID
+  `<task>
+  You are the final relevance judge in a hybrid search system. Your task is to re-rank search results by analyzing their true relevance to the user's query.
+  These documents have already passed through:
+  1. Semantic search (vector similarity)
+  2. Keyword-based search (full-text search)
  
-User query:
-{query}
+  Your ranking will determine the final order and which documents are shown to the user.
+</task>
+
+<context>
+  Each document contains:
+  - Main content text
+  - Optional summary/context
+  - Metadata (type, creation info)
+  The documents can be:
+  - Workspace nodes (various content types)
+  - Documents (files, notes)
+  - Database records
+</context>
+
+<ranking_criteria>
+  Evaluate relevance based on:
+  1. Direct answer presence (highest priority)
+     - Does the content directly answer the query?
+     - Are key details or facts present?
  
-Documents:
-{context}
+  2. Contextual relevance
+     - How well does the content relate to the query topic?
+     - Is the context/summary relevant?
+     - Does it provide important background information?
  
-Return an array of rankings in JSON format.`
+  3. Information freshness
+     - For time-sensitive queries, prefer recent content
+     - For conceptual queries, recency matters less
+  
+  4. Content completeness
+     - Does it provide comprehensive information?
+     - Are related concepts explained?
+  
+  5. Source appropriateness
+     - Is the document type appropriate for the query?
+     - Does the source authority match the information need?
+</ranking_criteria>
+
+<scoring_guidelines>
+  Score from 0 to 1, where:
+  1.0: Perfect match, directly answers query
+  0.8-0.9: Highly relevant, contains most key information
+  0.5-0.7: Moderately relevant, contains some useful information
+  0.2-0.4: Tangentially relevant, minimal useful information
+  0.0-0.1: Not relevant or useful for the query
+</scoring_guidelines>
+
+<documents>
+  {context}
+</documents>
+
+<user_query>
+  {query}
+</user_query>
+
+<output_format>
+  Return a JSON array of objects, each containing:
+  - "index": original position (integer)
+  - "score": relevance score (0-1 float)
+  - "type": document type (string)
+  - "sourceId": original source ID (string)
+  
+  Example:
+  [
+    {{"index": 2, "score": 0.95, "type": "document", "sourceId": "doc123"}},
+    {{"index": 0, "score": 0.7, "type": "node", "sourceId": "node456"}}
+  ]
+</output_format>`
 );

 export const answerPrompt = ChatPromptTemplate.fromTemplate(
-  `You are Colanode's AI assistant.
+  `<system_context>
+  You are an AI assistant in a collaboration workspace app called Colanode.

-CURRENT TIME: {currentTimestamp}
-WORKSPACE: {workspaceName}
-USER: {userName} ({userEmail})
+  CURRENT TIME: {currentTimestamp}
+  WORKSPACE: {workspaceName}
+  USER: {userName} ({userEmail})
+</system_context>

-CONVERSATION HISTORY:
-{formattedChatHistory}
+<current_conversation_history>
+  {formattedChatHistory}
+</current_conversation_history>

-RELEVANT CONTEXT:
-{formattedDocuments}
+<context>
+  {formattedDocuments}
+</context>

-USER QUERY:
-{question}
+<user_query>
+  {question}
+</user_query>

-Based solely on the conversation history and the relevant context above, provide a clear and professional answer to the user's query. In your answer, include exact quotes from the provided context that support your answer.
-If the relevant context does not contain any information that answers the user's query, respond with "No relevant information found."
+<task>
+  Based solely on the current conversation history and the relevant context above, provide a clear and professional answer to the user's query. In your answer, include exact quotes from the provided context that support your answer.
+  If the relevant context does not contain any information that answers the user's query, respond with "No relevant information found." This is a critical step to ensure correct answers.
+</task>

-Return your response as a JSON object with the following structure:
-{{
-  "answer": <your answer as a string>,
-  "citations": [
-    {{ "sourceId": <source id>, "quote": <exact quote from the context> }},
-    ...
-  ]
-}}`
+<output_format>
+  Return your response as a JSON object with the following structure:
+  {{
+    "answer": <your answer as a string>,
+    "citations": [
+      {{ "sourceId": <source id>, "quote": <exact quote from the context> }},
+      ...
+    ]
+  }}
+</output_format>`
 );

 export const intentRecognitionPrompt = PromptTemplate.fromTemplate(
-  `Determine if the following user query requires retrieving additional context.
-Return exactly one value: "retrieve" or "no_context".
+  `<task>
+  Determine if the following user query requires retrieving context from the workspace's knowledge base.
+  You are a crucial decision point in an AI assistant system that must decide between:
+  1. Retrieving and using specific context from the workspace ("retrieve")
+  2. Answering directly from general knowledge ("no_context")
+</task>

-Conversation History:
-{formattedChatHistory}
+<context>
+  This system has access to:
+  - Documents and their embeddings
+  - Node content (various types of workspace items)
+  - Database records and their fields
+  - Previous conversation history
+</context>

-User Query:
-{question}`
+<guidelines>
+  Return "retrieve" when the query:
+  - Asks about specific workspace content, documents, or data
+  - References previous conversations or shared content
+  - Mentions specific projects, tasks, or workspace items
+  - Requires up-to-date information from the workspace
+  - Contains temporal references to workspace activity
+  - Asks about specific people or collaborators
+  - Needs details about database records or fields
+
+  Return "no_context" when the query:
+  - Asks for general knowledge or common facts
+  - Requests simple calculations or conversions
+  - Asks about general concepts without workspace specifics
+  - Makes small talk
+  - Requests explanations of universal concepts
+  - Can be answered correctly without workspace-specific information
+</guidelines>
+
+<examples>
+  "retrieve" examples:
+  - "What did John say about the API design yesterday?"
+  - "Show me the latest documentation about user authentication"
+  - "Find records in the Projects database where status is completed"
+  - "What were the key points from our last meeting?"
+
+  "no_context" examples:
+  - "What is REST API?"
+  - "How do I write a good commit message?"
+  - "Convert 42 kilometers to miles"
+  - "What's your name?"
+  - "Explain what is Docker in simple terms"
+</examples>
+
+<conversation_history>
+  {formattedChatHistory}
+</conversation_history>
+
+<user_query>
+  {question}
+</user_query>
+
+<output_format>
+  Return exactly one value: "retrieve" or "no_context"
+</output_format>`
 );

 export const noContextPrompt = PromptTemplate.fromTemplate(
-  `Answer the following query concisely using general knowledge, without retrieving additional context.
+  `<task>
+  Answer the following query concisely using general knowledge, without retrieving additional context. Return only the answer.
+</task>

-Conversation History:
-{formattedChatHistory}
+<conversation_history>
+  {formattedChatHistory}
+</conversation_history>

-User Query:
-{question}
-
-Return only the answer.`
+<user_query>
+  {question}
+</user_query>`
 );

 export const databaseFilterPrompt = ChatPromptTemplate.fromTemplate(
--- a/apps/server/src/services/document-retrieval-service.ts
+++ b/apps/server/src/services/document-retrieval-service.ts
@@ -4,6 +4,8 @@ import { database } from '@/data/database';
 import { configuration } from '@/lib/configuration';
 import { sql } from 'kysely';
 import { SearchResult } from '@/types/retrieval';
+import { RewrittenQuery } from '@/types/llm';
+
 export class DocumentRetrievalService {
  private embeddings = new OpenAIEmbeddings({
    apiKey: configuration.ai.embedding.apiKey,
@@ -12,13 +14,15 @@ export class DocumentRetrievalService {
  });

  public async retrieve(
-    query: string,
+    rewrittenQuery: RewrittenQuery,
    workspaceId: string,
    userId: string,
    limit = configuration.ai.retrieval.hybridSearch.maxResults,
    contextNodeIds?: string[]
  ): Promise<Document[]> {
-    const embedding = await this.embeddings.embedQuery(query);
+    const embedding = await this.embeddings.embedQuery(
+      rewrittenQuery.semanticQuery
+    );
    if (!embedding) {
      return [];
    }
@@ -31,7 +35,13 @@ export class DocumentRetrievalService {
        limit,
        contextNodeIds
      ),
-      this.keywordSearch(query, workspaceId, userId, limit, contextNodeIds),
+      this.keywordSearch(
+        rewrittenQuery.keywordQuery,
+        workspaceId,
+        userId,
+        limit,
+        contextNodeIds
+      ),
    ]);

    return this.combineSearchResults(semanticResults, keywordResults);
@@ -57,6 +67,7 @@ export class DocumentRetrievalService {
      .select((eb) => [
        'document_embeddings.document_id as id',
        'document_embeddings.text',
+        'document_embeddings.summary',
        'documents.created_at',
        'documents.created_by',
        'document_embeddings.chunk as chunk_index',
@@ -78,6 +89,7 @@ export class DocumentRetrievalService {
      .groupBy([
        'document_embeddings.document_id',
        'document_embeddings.text',
+        'document_embeddings.summary',
        'documents.created_at',
        'documents.created_by',
        'document_embeddings.chunk',
@@ -89,6 +101,7 @@ export class DocumentRetrievalService {
    return results.map((result) => ({
      id: result.id,
      text: result.text,
+      summary: result.summary,
      score: result.similarity,
      type: 'semantic',
      createdAt: result.created_at,
@@ -117,6 +130,7 @@ export class DocumentRetrievalService {
      .select((eb) => [
        'document_embeddings.document_id as id',
        'document_embeddings.text',
+        'document_embeddings.summary',
        'documents.created_at',
        'documents.created_by',
        'document_embeddings.chunk as chunk_index',
@@ -145,6 +159,7 @@ export class DocumentRetrievalService {
        'documents.created_at',
        'documents.created_by',
        'document_embeddings.chunk',
+        'document_embeddings.summary',
      ])
      .orderBy('rank', 'desc')
      .limit(limit)
@@ -153,6 +168,7 @@ export class DocumentRetrievalService {
    return results.map((result) => ({
      id: result.id,
      text: result.text,
+      summary: result.summary,
      score: result.rank,
      type: 'keyword',
      createdAt: result.created_at,
@@ -243,7 +259,7 @@ export class DocumentRetrievalService {
          ? authorMap.get(result.createdBy)
          : null;
        return new Document({
-          pageContent: result.text,
+          pageContent: `${result.summary}\n\n${result.text}`,
          metadata: {
            id: result.id,
            score: result.finalScore,
@@ -253,7 +269,7 @@ export class DocumentRetrievalService {
            author: author
              ? {
                  id: author.id,
-                  name: author.name || 'Anonymous',
+                  name: author.name || 'Unknown',
                }
              : null,
          },
--- a/apps/server/src/services/llm-service.ts
+++ b/apps/server/src/services/llm-service.ts
@@ -11,6 +11,8 @@ import {
  CitedAnswer,
  databaseFilterSchema,
  DatabaseFilterResult,
+  RewrittenQuery,
+  rewrittenQuerySchema,
 } from '@/types/llm';

 import {
@@ -61,13 +63,10 @@ const getChatModel = (
  }
 };

-export const rewriteQuery = async (query: string): Promise<string> => {
+export const rewriteQuery = async (query: string): Promise<RewrittenQuery> => {
  const task = 'queryRewrite';
-  const model = getChatModel(task);
-  return queryRewritePrompt
-    .pipe(model)
-    .pipe(new StringOutputParser())
-    .invoke({ query });
+  const model = getChatModel(task).withStructuredOutput(rewrittenQuerySchema);
+  return queryRewritePrompt.pipe(model).invoke({ query });
 };

 export const summarizeDocument = async (
--- a/apps/server/src/services/node-retrieval-service.ts
+++ b/apps/server/src/services/node-retrieval-service.ts
@@ -4,7 +4,7 @@ import { database } from '@/data/database';
 import { configuration } from '@/lib/configuration';
 import { sql } from 'kysely';
 import { SearchResult } from '@/types/retrieval';
-
+import { RewrittenQuery } from '@/types/llm';
 export class NodeRetrievalService {
  private embeddings = new OpenAIEmbeddings({
    apiKey: configuration.ai.embedding.apiKey,
@@ -13,13 +13,15 @@ export class NodeRetrievalService {
  });

  public async retrieve(
-    query: string,
+    rewrittenQuery: RewrittenQuery,
    workspaceId: string,
    userId: string,
    limit = configuration.ai.retrieval.hybridSearch.maxResults,
    contextNodeIds?: string[]
  ): Promise<Document[]> {
-    const embedding = await this.embeddings.embedQuery(query);
+    const embedding = await this.embeddings.embedQuery(
+      rewrittenQuery.semanticQuery
+    );
    if (!embedding) {
      return [];
    }
@@ -32,7 +34,13 @@ export class NodeRetrievalService {
        limit,
        contextNodeIds
      ),
-      this.keywordSearch(query, workspaceId, userId, limit, contextNodeIds),
+      this.keywordSearch(
+        rewrittenQuery.keywordQuery,
+        workspaceId,
+        userId,
+        limit,
+        contextNodeIds
+      ),
    ]);

    return this.combineSearchResults(semanticResults, keywordResults);
@@ -57,6 +65,7 @@ export class NodeRetrievalService {
      .select((eb) => [
        'node_embeddings.node_id as id',
        'node_embeddings.text',
+        'node_embeddings.summary',
        'nodes.created_at',
        'nodes.created_by',
        'node_embeddings.chunk as chunk_index',
@@ -81,6 +90,7 @@ export class NodeRetrievalService {
        'nodes.created_at',
        'nodes.created_by',
        'node_embeddings.chunk',
+        'node_embeddings.summary',
      ])
      .orderBy('similarity', 'asc')
      .limit(limit)
@@ -89,6 +99,7 @@ export class NodeRetrievalService {
    return results.map((result) => ({
      id: result.id,
      text: result.text,
+      summary: result.summary,
      score: result.similarity,
      type: 'semantic',
      createdAt: result.created_at,
@@ -116,6 +127,7 @@ export class NodeRetrievalService {
      .select((eb) => [
        'node_embeddings.node_id as id',
        'node_embeddings.text',
+        'node_embeddings.summary',
        'nodes.created_at',
        'nodes.created_by',
        'node_embeddings.chunk as chunk_index',
@@ -144,6 +156,7 @@ export class NodeRetrievalService {
        'nodes.created_at',
        'nodes.created_by',
        'node_embeddings.chunk',
+        'node_embeddings.summary',
      ])
      .orderBy('rank', 'desc')
      .limit(limit)
@@ -152,6 +165,7 @@ export class NodeRetrievalService {
    return results.map((result) => ({
      id: result.id,
      text: result.text,
+      summary: result.summary,
      score: result.rank,
      type: 'keyword',
      createdAt: result.created_at,
@@ -242,7 +256,7 @@ export class NodeRetrievalService {
          ? authorMap.get(result.createdBy)
          : null;
        return new Document({
-          pageContent: result.text,
+          pageContent: `${result.summary}\n\n${result.text}`,
          metadata: {
            id: result.id,
            score: result.finalScore,
@@ -252,7 +266,7 @@ export class NodeRetrievalService {
            author: author
              ? {
                  id: author.id,
-                  name: author.name || 'Anonymous',
+                  name: author.name || 'Unknown',
                }
              : null,
          },
--- a/apps/server/src/types/assistant.ts
+++ b/apps/server/src/types/assistant.ts
@@ -1,5 +1,11 @@
 import { Document } from '@langchain/core/documents';
 import { Annotation } from '@langchain/langgraph';
+import {
+  RerankedDocuments,
+  CitedAnswer,
+  DatabaseFilterResult,
+  RewrittenQuery,
+} from './llm';

 export type Citation = {
  sourceId: string;
@@ -30,20 +36,13 @@ export const ResponseState = Annotation.Root({
  userDetails: Annotation<{ name: string; email: string }>(),
  parentMessageId: Annotation<string>(),
  currentMessageId: Annotation<string>(),
-  rewrittenQuery: Annotation<string>(),
+  rewrittenQuery: Annotation<RewrittenQuery>(),
  contextDocuments: Annotation<Document[]>(),
  chatHistory: Annotation<Document[]>(),
-  rerankedContext: Annotation<
-    Array<{
-      index: number;
-      score: number;
-      type: string;
-      sourceId: string;
-    }>
-  >(),
+  rerankedContext: Annotation<RerankedDocuments['rankings']>(),
  topContext: Annotation<Document[]>(),
  finalAnswer: Annotation<string>(),
-  citations: Annotation<Array<{ sourceId: string; quote: string }>>(),
+  citations: Annotation<CitedAnswer['citations']>(),
  originalMessage: Annotation<any>(),
  intent: Annotation<'retrieve' | 'no_context'>(),
  databaseContext: Annotation<
@@ -54,13 +53,7 @@ export const ResponseState = Annotation.Root({
      sampleRecords: any[];
    }>
  >(),
-  databaseFilters: Annotation<{
-    shouldFilter: boolean;
-    filters: Array<{
-      databaseId: string;
-      filters: any[];
-    }>;
-  }>(),
+  databaseFilters: Annotation<DatabaseFilterResult>(),
  selectedContextNodeIds: Annotation<string[]>(),
 });

--- a/apps/server/src/types/llm.ts
+++ b/apps/server/src/types/llm.ts
@@ -33,3 +33,10 @@ export const databaseFilterSchema = z.object({
  ),
 });
 export type DatabaseFilterResult = z.infer<typeof databaseFilterSchema>;
+
+export const rewrittenQuerySchema = z.object({
+  semanticQuery: z.string(),
+  keywordQuery: z.string(),
+});
+
+export type RewrittenQuery = z.infer<typeof rewrittenQuerySchema>;
--- a/apps/server/src/types/retrieval.ts
+++ b/apps/server/src/types/retrieval.ts
@@ -1,6 +1,7 @@
 export type SearchResult = {
  id: string;
  text: string;
+  summary: string | null;
  score: number;
  type: 'semantic' | 'keyword';
  createdAt?: Date;