Skip to content

LCORE-383: Removing attachment metadata at conversation end point and streamlining o… #983

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sriroopar wants to merge 1 commit into
base: main
Choose a base branch
from
Open

LCORE-383: Removing attachment metadata at conversation end point and streamlining o… #983

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/app/endpoints/conversations_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -257,7 +257,7 @@ def build_conversation_turn_from_cache_entry(entry: CacheEntry) -> ConversationT
"""
# Create Message objects for user and assistant
messages = [
Message(content=entry.query, type="user"),
Message(content=entry.query, type="user", attachments=entry.attachments),
Message(content=entry.response, type="assistant"),
]

Expand Down
42 changes: 39 additions & 3 deletions src/cache/postgres_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ class PostgresCache(Cache):
referenced_documents | jsonb | |
tool_calls | jsonb | |
tool_results | jsonb | |
attachments | jsonb | |
Indexes:
"cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
"timestamps" btree (created_at)
Expand All @@ -60,6 +61,7 @@ class PostgresCache(Cache):
referenced_documents jsonb,
tool_calls jsonb,
tool_results jsonb,
attachments jsonb,
PRIMARY KEY(user_id, conversation_id, created_at)
);
"""
Expand All @@ -81,7 +83,7 @@ class PostgresCache(Cache):

SELECT_CONVERSATION_HISTORY_STATEMENT = """
SELECT query, response, provider, model, started_at, completed_at,
referenced_documents, tool_calls, tool_results
referenced_documents, tool_calls, tool_results, attachments
FROM cache
WHERE user_id=%s AND conversation_id=%s
ORDER BY created_at
Expand All @@ -90,8 +92,8 @@ class PostgresCache(Cache):
INSERT_CONVERSATION_HISTORY_STATEMENT = """
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
query, response, provider, model, referenced_documents,
tool_calls, tool_results)
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s, %s, %s)
tool_calls, tool_results, attachments)
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""

QUERY_CACHE_SIZE = """
Expand Down Expand Up @@ -321,6 +323,24 @@ def get( # pylint: disable=R0914
e,
)

# Parse attachments back into Attachment objects
attachments_data = conversation_entry[9]
attachments_obj = None
if attachments_data:
try:
from models.requests import Attachment

attachments_obj = [
Attachment.model_validate(att) for att in attachments_data
]
except (ValueError, TypeError) as e:
logger.warning(
"Failed to deserialize attachments for "
"conversation %s: %s",
conversation_id,
e,
)

cache_entry = CacheEntry(
query=conversation_entry[0],
response=conversation_entry[1],
Expand All @@ -331,6 +351,7 @@ def get( # pylint: disable=R0914
referenced_documents=docs_obj,
tool_calls=tool_calls_obj,
tool_results=tool_results_obj,
attachments=attachments_obj,
)
result.append(cache_entry)

Expand Down Expand Up @@ -405,6 +426,20 @@ def insert_or_append(
e,
)

attachments_json = None
if cache_entry.attachments:
try:
attachments_as_dicts = [
att.model_dump(mode="json") for att in cache_entry.attachments
]
attachments_json = json.dumps(attachments_as_dicts)
except (TypeError, ValueError) as e:
logger.warning(
"Failed to serialize attachments for conversation %s: %s",
conversation_id,
e,
)

# the whole operation is run in one transaction
with self.connection.cursor() as cursor:
cursor.execute(
Expand All @@ -421,6 +456,7 @@ def insert_or_append(
referenced_documents_json,
tool_calls_json,
tool_results_json,
attachments_json,
),
)

Expand Down
42 changes: 39 additions & 3 deletions src/cache/sqlite_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class SQLiteCache(Cache):
referenced_documents | text | |
tool_calls | text | |
tool_results | text | |
attachments | text | |
Indexes:
"cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
"cache_key_key" UNIQUE CONSTRAINT, btree (key)
Expand All @@ -59,6 +60,7 @@ class SQLiteCache(Cache):
referenced_documents text,
tool_calls text,
tool_results text,
attachments text,
PRIMARY KEY(user_id, conversation_id, created_at)
);
"""
Expand All @@ -80,7 +82,7 @@ class SQLiteCache(Cache):

SELECT_CONVERSATION_HISTORY_STATEMENT = """
SELECT query, response, provider, model, started_at, completed_at,
referenced_documents, tool_calls, tool_results
referenced_documents, tool_calls, tool_results, attachments
FROM cache
WHERE user_id=? AND conversation_id=?
ORDER BY created_at
Expand All @@ -89,8 +91,8 @@ class SQLiteCache(Cache):
INSERT_CONVERSATION_HISTORY_STATEMENT = """
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
query, response, provider, model, referenced_documents,
tool_calls, tool_results)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
tool_calls, tool_results, attachments)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"""

QUERY_CACHE_SIZE = """
Expand Down Expand Up @@ -301,6 +303,24 @@ def get( # pylint: disable=R0914
e,
)

# Parse attachments back into Attachment objects
attachments_json_str = conversation_entry[9]
attachments_obj = None
if attachments_json_str:
try:
from models.requests import Attachment

attachments_data = json.loads(attachments_json_str)
attachments_obj = [
Attachment.model_validate(att) for att in attachments_data
]
except (json.JSONDecodeError, ValueError) as e:
logger.warning(
"Failed to deserialize attachments for conversation %s: %s",
conversation_id,
e,
)

cache_entry = CacheEntry(
query=conversation_entry[0],
response=conversation_entry[1],
Expand All @@ -311,6 +331,7 @@ def get( # pylint: disable=R0914
referenced_documents=docs_obj,
tool_calls=tool_calls_obj,
tool_results=tool_results_obj,
attachments=attachments_obj,
)
result.append(cache_entry)

Expand Down Expand Up @@ -386,6 +407,20 @@ def insert_or_append(
e,
)

attachments_json = None
if cache_entry.attachments:
try:
attachments_as_dicts = [
att.model_dump(mode="json") for att in cache_entry.attachments
]
attachments_json = json.dumps(attachments_as_dicts)
except (TypeError, ValueError) as e:
logger.warning(
"Failed to serialize attachments for conversation %s: %s",
conversation_id,
e,
)

cursor.execute(
self.INSERT_CONVERSATION_HISTORY_STATEMENT,
(
Expand All @@ -401,6 +436,7 @@ def insert_or_append(
referenced_documents_json,
tool_calls_json,
tool_results_json,
attachments_json,
),
)

Expand Down
7 changes: 6 additions & 1 deletion src/models/cache_entry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
"""Model for conversation history cache entry."""

from typing import Optional
from typing import TYPE_CHECKING, Optional
from pydantic import BaseModel
from utils.types import ReferencedDocument, ToolCallSummary, ToolResultSummary

if TYPE_CHECKING:
from models.requests import Attachment


class CacheEntry(BaseModel):
"""Model representing a cache entry.
Expand All @@ -16,6 +19,7 @@ class CacheEntry(BaseModel):
referenced_documents: List of documents referenced by the response
tool_calls: List of tool calls made during response generation
tool_results: List of tool results from tool calls
attachments: Optional list of attachments included with the query
"""

query: str
Expand All @@ -27,3 +31,4 @@ class CacheEntry(BaseModel):
referenced_documents: Optional[list[ReferencedDocument]] = None
tool_calls: Optional[list[ToolCallSummary]] = None
tool_results: Optional[list[ToolResultSummary]] = None
attachments: Optional[list["Attachment"]] = None
10 changes: 9 additions & 1 deletion src/models/responses.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

"""Models for REST API responses."""

from typing import Any, ClassVar, Literal, Optional, Union
from typing import Any, ClassVar, Literal, Optional, Union, TYPE_CHECKING

from fastapi import status
from pydantic import BaseModel, Field
Expand All @@ -13,6 +13,9 @@
from quota.quota_exceed_error import QuotaExceedError
from utils.types import RAGChunk, ReferencedDocument, ToolCallSummary, ToolResultSummary

if TYPE_CHECKING:
from models.requests import Attachment

Comment on lines +16 to +18
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify models/requests.py does NOT import from models/responses.py (circular import check)
echo "=== Imports in models/requests.py ==="
rg -n "from models.responses\|import models.responses" --type py

# Verify no model_rebuild() call exists for Message
echo ""
echo "=== model_rebuild calls ==="
rg -n "model_rebuild" --type py src/

Repository: lightspeed-core/lightspeed-stack

Length of output: 142

🏁 Script executed:

#!/bin/bash
echo "=== Lines 1-30 of src/models/responses.py ==="
head -30 src/models/responses.py | cat -n

echo ""
echo "=== Lines 825-835 of src/models/responses.py ==="
sed -n '825,835p' src/models/responses.py | cat -n

echo ""
echo "=== Check if Attachment is used in responses.py ==="
rg -n "Attachment" src/models/responses.py

Repository: lightspeed-core/lightspeed-stack

Length of output: 2073

🏁 Script executed:

#!/bin/bash
echo "=== All TYPE_CHECKING imports in responses.py ==="
rg -n "if TYPE_CHECKING" src/models/responses.py -A 2

echo ""
echo "=== Verify Message class definition ==="
rg -n "^class Message" src/models/responses.py -A 5

Repository: lightspeed-core/lightspeed-stack

Length of output: 437

Remove TYPE_CHECKING and directly import Attachment to fix runtime PydanticUndefinedAnnotation error.

Attachment is imported only inside if TYPE_CHECKING:, so it is never in models.responses's global namespace at runtime. Pydantic v2 resolves the string forward-reference "Attachment" by looking it up in the defining module's __dict__, so calling Message.model_json_schema() (which FastAPI does at startup to generate OpenAPI docs) will raise PydanticUndefinedAnnotation.

Verification confirms no circular import exists; models/requests.py does not import from models/responses.py. The fix is to directly import Attachment at module level and update the field annotation:

Fix 
-from typing import Any, ClassVar, Literal, Optional, Union, TYPE_CHECKING
+from typing import Any, ClassVar, Literal, Optional, Union

 from fastapi import status
 from pydantic import BaseModel, Field
 from pydantic_core import SchemaError

 from constants import MEDIA_TYPE_EVENT_STREAM
 from models.config import Action, Configuration
+from models.requests import Attachment
 from quota.quota_exceed_error import QuotaExceedError
 from utils.types import RAGChunk, ReferencedDocument, ToolCallSummary, ToolResultSummary

-if TYPE_CHECKING:
-    from models.requests import Attachment
-

Update line 829:

-    attachments: Optional[list["Attachment"]] = Field(
+    attachments: Optional[list[Attachment]] = Field(
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/models/responses.py` around lines 16 - 18, The runtime
PydanticUndefinedAnnotation occurs because Attachment is only imported inside
the TYPE_CHECKING block so it's not present when Message.model_json_schema() is
called; remove the if TYPE_CHECKING guard and import Attachment at module level
in models.responses, then change any string forward-reference annotations (e.g.,
"Attachment") to the actual Attachment type on the Message (or related) field so
Pydantic can resolve it at runtime (this fixes the error raised when calling
Message.model_json_schema()).

SUCCESSFUL_RESPONSE_DESCRIPTION = "Successful response"
BAD_REQUEST_DESCRIPTION = "Invalid request format"
UNAUTHORIZED_DESCRIPTION = "Unauthorized"
Expand Down Expand Up @@ -810,6 +813,7 @@ class Message(BaseModel):
Attributes:
content: The message content.
type: The type of message.
attachments: Optional list of attachments included with the message.
"""

content: str = Field(
Expand All @@ -822,6 +826,10 @@ class Message(BaseModel):
description="The type of message",
examples=["user", "assistant", "system", "developer"],
)
attachments: Optional[list["Attachment"]] = Field(
default=None,
description="Optional attachments included with this message",
)


class ConversationTurn(BaseModel):
Expand Down
1 change: 1 addition & 0 deletions src/utils/query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,7 @@ def store_query_results( # pylint: disable=too-many-arguments,too-many-locals
referenced_documents=summary.referenced_documents,
tool_calls=summary.tool_calls,
tool_results=summary.tool_results,
attachments=query_request.attachments,
)

logger.info("Storing conversation in cache")
Expand Down
Loading