feat: add conversation id header to stored interactions (#91)

enitrat · web-flow · commit 7d9ae48b1ac0 · 2025-12-01T20:40:13.000+01:00
diff --git a/python/pyproject.toml b/python/pyproject.toml
@@ -43,8 +43,7 @@ dependencies = [
   "pydantic-settings>=2.0.0",
   "pydantic>=2.0.0",
   "pyinstrument>=5.1.1",
-  "pytest-asyncio>=1.0.0",
-  "pytest>=8.4.1",
+
   "python-dotenv>=1.0.0",
   "python-multipart>=0.0.6",
   "structlog>=24.0.0",
@@ -55,20 +54,13 @@ dependencies = [
   "xai_sdk>=1.3.1",
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
-  "pytest>=8.0.0",
-  "pytest-asyncio>=0.23.0",
-  "pytest-cov>=5.0.0",
-  "pytest-benchmark>=4.0.0",
-  "pytest-mock>=3.0.0",
-  "black>=24.0.0",
-  "ruff>=0.4.0",
-  "mypy>=1.0.0",
-  "types-toml>=0.10.0",
-  "pre-commit>=3.0.0",
-  "testcontainers[postgres]>=4.0.0",
   "nest-asyncio>=1.6.0",
+  "ty>=0.0.1a15",
+  "testcontainers[postgres]>=4.13.0",
+  "pytest-asyncio>=1.0.0",
+  "pytest>=8.4.1",
 ]
 
 [project.scripts]
@@ -163,6 +155,3 @@ exclude_lines = [
   "raise NotImplementedError",
   "if TYPE_CHECKING:",
 ]
-
-[dependency-groups]
-dev = ["nest-asyncio>=1.6.0", "ty>=0.0.1a15"]
diff --git a/python/src/cairo_coder/db/models.py b/python/src/cairo_coder/db/models.py
@@ -20,6 +20,7 @@ class UserInteraction(BaseModel):
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
     agent_id: str
     mcp_mode: bool = False
+    conversation_id: Optional[str] = None
     chat_history: Optional[list[dict[str, Any]]] = None
     query: str
     generated_answer: Optional[str] = None
diff --git a/python/src/cairo_coder/db/repository.py b/python/src/cairo_coder/db/repository.py
@@ -83,17 +83,19 @@ async def create_user_interaction(interaction: UserInteraction) -> None:
                     id,
                     agent_id,
                     mcp_mode,
+                    conversation_id,
                     chat_history,
                     query,
                     generated_answer,
                     retrieved_sources,
                     llm_usage
                 )
-                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                 """,
                 interaction.id,
                 interaction.agent_id,
                 interaction.mcp_mode,
+                interaction.conversation_id,
                 _serialize_json_field(interaction.chat_history),
                 interaction.query,
                 interaction.generated_answer,
@@ -112,6 +114,7 @@ async def get_interactions(
     limit: int,
     offset: int,
     query_text: str | None = None,
+    conversation_id: str | None = None,
 ) -> tuple[list[dict[str, Any]], int]:
     """Fetch paginated interactions matching the supplied filters.
 
@@ -139,6 +142,10 @@ async def get_interactions(
             params.append(f"%{query_text}%")
             filters.append(f"query ILIKE ${len(params)}")
 
+        if conversation_id:
+            params.append(conversation_id)
+            filters.append(f"conversation_id = ${len(params)}")
+
         where_clause = "WHERE " + " AND ".join(filters) if filters else ""
 
         count_query = f"""
@@ -152,7 +159,7 @@ async def get_interactions(
         limit_placeholder = len(params) - 1
         offset_placeholder = len(params)
         data_query = f"""
-            SELECT id, created_at, agent_id, query, chat_history, generated_answer
+            SELECT id, created_at, agent_id, query, chat_history, generated_answer, conversation_id
             FROM user_interactions
             {where_clause}
             ORDER BY created_at DESC
@@ -192,17 +199,19 @@ async def migrate_user_interaction(interaction: UserInteraction) -> tuple[bool,
                     created_at,
                     agent_id,
                     mcp_mode,
+                    conversation_id,
                     chat_history,
                     query,
                     generated_answer,
                     retrieved_sources,
                     llm_usage
                 )
-                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
                 ON CONFLICT (id) DO UPDATE SET
                     created_at = EXCLUDED.created_at,
                     agent_id = EXCLUDED.agent_id,
                     mcp_mode = EXCLUDED.mcp_mode,
+                    conversation_id = EXCLUDED.conversation_id,
                     chat_history = EXCLUDED.chat_history,
                     query = EXCLUDED.query,
                     generated_answer = EXCLUDED.generated_answer,
@@ -214,6 +223,7 @@ async def migrate_user_interaction(interaction: UserInteraction) -> tuple[bool,
                 interaction.created_at,
                 interaction.agent_id,
                 interaction.mcp_mode,
+                interaction.conversation_id,
                 _serialize_json_field(interaction.chat_history),
                 interaction.query,
                 interaction.generated_answer,
diff --git a/python/src/cairo_coder/db/session.py b/python/src/cairo_coder/db/session.py
@@ -71,6 +71,7 @@ async def execute_schema_scripts() -> None:
                 created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                 agent_id VARCHAR(50) NOT NULL,
                 mcp_mode BOOLEAN NOT NULL DEFAULT FALSE,
+                conversation_id VARCHAR(100),
                 chat_history JSONB,
                 query TEXT NOT NULL,
                 generated_answer TEXT,
@@ -79,12 +80,21 @@ async def execute_schema_scripts() -> None:
             );
             """
         )
+        # Migration: add conversation_id column if it doesn't exist (for existing tables)
+        await connection.execute(
+            """
+            ALTER TABLE user_interactions
+            ADD COLUMN IF NOT EXISTS conversation_id VARCHAR(100);
+            """
+        )
         await connection.execute(
             """
             CREATE INDEX IF NOT EXISTS idx_interactions_created_at
                 ON user_interactions(created_at);
             CREATE INDEX IF NOT EXISTS idx_interactions_agent_id
                 ON user_interactions(agent_id);
+            CREATE INDEX IF NOT EXISTS idx_interactions_conversation_id
+                ON user_interactions(conversation_id);
             """
         )
     logger.info("Database schema initialized.")
diff --git a/python/src/cairo_coder/server/app.py b/python/src/cairo_coder/server/app.py
@@ -151,6 +151,7 @@ async def log_interaction_task(
     chat_history: list[Message],
     response: ChatCompletionResponse,
     agent: RagPipeline,
+    conversation_id: str | None = None,
 ) -> None:
     """Background task that persists a user interaction."""
     sources_data = [
@@ -167,6 +168,7 @@ async def log_interaction_task(
     interaction = UserInteraction(
         agent_id=agent_id,
         mcp_mode=mcp_mode,
+        conversation_id=conversation_id,
         chat_history=chat_history_dicts,
         query=query,
         generated_answer=response.choices[0].message.content if response.choices else None,
@@ -183,6 +185,7 @@ async def log_interaction_raw(
     chat_history: list[Message],
     generated_answer: str | None,
     agent: RagPipeline,
+    conversation_id: str | None = None,
 ) -> None:
     """Persist a user interaction without constructing a full response object."""
     sources_data = [
@@ -198,6 +201,7 @@ async def log_interaction_raw(
     interaction = UserInteraction(
         agent_id=agent_id,
         mcp_mode=mcp_mode,
+        conversation_id=conversation_id,
         chat_history=chat_history_dicts,
         query=query,
         generated_answer=generated_answer,
@@ -423,6 +427,9 @@ async def _handle_chat_completion(
         vector_db: SourceFilteredPgVectorRM | None = None,
     ):
         """Handle chat completion request."""
+        # Extract conversation ID from header
+        conversation_id = req.headers.get("x-conversation-id")
+
         # Convert messages to internal format
         messages = []
         for msg in request.messages:
@@ -443,7 +450,9 @@ async def _handle_chat_completion(
         # Handle streaming vs non-streaming
         if request.stream:
             return StreamingResponse(
-                self._stream_chat_completion(agent, query, messages[:-1], mcp_mode, effective_agent_id),
+                self._stream_chat_completion(
+                    agent, query, messages[:-1], mcp_mode, effective_agent_id, conversation_id
+                ),
                 media_type="text/event-stream",
                 headers={
                     "Cache-Control": "no-cache",
@@ -462,12 +471,19 @@ async def _handle_chat_completion(
             chat_history=chat_history,
             response=response,
             agent=agent,
+            conversation_id=conversation_id,
         )
 
         return response
 
     async def _stream_chat_completion(
-        self, agent: RagPipeline, query: str, history: list[Message], mcp_mode: bool, agent_id: str
+        self,
+        agent: RagPipeline,
+        query: str,
+        history: list[Message],
+        mcp_mode: bool,
+        agent_id: str,
+        conversation_id: str | None = None,
     ) -> AsyncGenerator[str, None]:
         """Stream chat completion response - replicates TypeScript streaming."""
         response_id = str(uuid.uuid4())
@@ -580,6 +596,7 @@ async def _stream_chat_completion(
                     chat_history=history,
                     generated_answer=final_response,
                     agent=agent,
+                    conversation_id=conversation_id,
                 )
             except Exception as log_error:
                 logger.error("Failed to log streaming interaction", error=str(log_error), exc_info=True)
diff --git a/python/src/cairo_coder/server/insights_api.py b/python/src/cairo_coder/server/insights_api.py
@@ -28,6 +28,7 @@ class QueryResponse(BaseModel):
     query: str
     chat_history: list[dict[str, Any]]
     output: str | None
+    conversation_id: str | None = None
 
 
 class PaginatedQueryResponse(BaseModel):
@@ -45,16 +46,19 @@ async def get_raw_queries(
     end_date: datetime | None = None,
     agent_id: str | None = None,
     query_text: str | None = None,
+    conversation_id: str | None = None,
     limit: int = 100,
     offset: int = 0,
 ) -> PaginatedQueryResponse:
     """Return raw user queries.
 
     If start_date and end_date are not provided, returns the last N queries
     ordered by creation time (where N is the limit parameter).
+
+    Use conversation_id to filter queries belonging to a specific conversation.
     """
     items, total = await get_interactions(
-        start_date, end_date, agent_id, limit, offset, query_text
+        start_date, end_date, agent_id, limit, offset, query_text, conversation_id
     )
     # Map generated_answer to output for API response
     responses = [
@@ -65,6 +69,7 @@ async def get_raw_queries(
             query=item["query"],
             chat_history=item["chat_history"] or [],
             output=item.get("generated_answer"),
+            conversation_id=item.get("conversation_id"),
         )
         for item in items
     ]
diff --git a/python/tests/integration/conftest.py b/python/tests/integration/conftest.py
@@ -182,6 +182,7 @@ async def test_db_pool(postgres_container):
                 created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                 agent_id VARCHAR(50) NOT NULL,
                 mcp_mode BOOLEAN NOT NULL DEFAULT FALSE,
+                conversation_id VARCHAR(100),
                 chat_history JSONB,
                 query TEXT NOT NULL,
                 generated_answer TEXT,
@@ -196,6 +197,8 @@ async def test_db_pool(postgres_container):
                 ON user_interactions(created_at);
             CREATE INDEX IF NOT EXISTS idx_interactions_agent_id
                 ON user_interactions(agent_id);
+            CREATE INDEX IF NOT EXISTS idx_interactions_conversation_id
+                ON user_interactions(conversation_id);
             """
         )
 
diff --git a/python/tests/integration/test_insights_api.py b/python/tests/integration/test_insights_api.py
@@ -155,6 +155,72 @@ def test_get_queries_without_dates_with_filters(self, client, populated_db_conne
         assert data["total"] >= 1
         assert all(item["agent_id"] == "cairo-coder" for item in data["items"])
 
+    def test_get_queries_with_conversation_id_filter(self, client, db_connection):
+        """Test that queries can be filtered by conversation_id."""
+        import asyncio
+        import json as _json
+        import uuid
+        from datetime import datetime, timedelta, timezone
+
+        now = datetime.now(timezone.utc)
+        conv_id = "test-conversation-123"
+
+        # Seed records with and without conversation_id
+        async def seed():
+            await db_connection.execute(
+                """
+                INSERT INTO user_interactions (id, created_at, agent_id, mcp_mode, conversation_id, chat_history, query, generated_answer)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8),
+                       ($9, $10, $11, $12, $13, $14, $15, $16),
+                       ($17, $18, $19, $20, $21, $22, $23, $24)
+                """,
+                uuid.uuid4(), now - timedelta(hours=2), "cairo-coder", False, conv_id, _json.dumps([]), "First msg", "Response 1",
+                uuid.uuid4(), now - timedelta(hours=1), "cairo-coder", False, conv_id, _json.dumps([{"role": "user", "content": "First msg"}]), "Second msg", "Response 2",
+                uuid.uuid4(), now - timedelta(minutes=30), "cairo-coder", False, None, _json.dumps([]), "Other msg", "Other response",
+            )
+
+        asyncio.get_event_loop().run_until_complete(seed())
+
+        # Filter by conversation_id
+        resp = client.get(
+            "/v1/insights/queries",
+            params={"conversation_id": conv_id, "limit": 100, "offset": 0},
+        )
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["total"] == 2
+        assert all(item["conversation_id"] == conv_id for item in data["items"])
+
+    def test_get_queries_returns_conversation_id(self, client, db_connection):
+        """Test that conversation_id is included in the response."""
+        import asyncio
+        import json as _json
+        import uuid
+        from datetime import datetime, timezone
+
+        now = datetime.now(timezone.utc)
+        conv_id = "response-test-conv-456"
+
+        async def seed():
+            await db_connection.execute(
+                """
+                INSERT INTO user_interactions (id, created_at, agent_id, mcp_mode, conversation_id, chat_history, query, generated_answer)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+                """,
+                uuid.uuid4(), now, "cairo-coder", False, conv_id, _json.dumps([]), "Test query", "Test response",
+            )
+
+        asyncio.get_event_loop().run_until_complete(seed())
+
+        resp = client.get("/v1/insights/queries", params={"limit": 100, "offset": 0})
+        assert resp.status_code == 200
+        data = resp.json()
+
+        # Find our seeded record and verify conversation_id is present
+        matching = [item for item in data["items"] if item.get("conversation_id") == conv_id]
+        assert len(matching) == 1
+        assert matching[0]["conversation_id"] == conv_id
+
 
 class TestDataIngestion:
     async def test_chat_completion_logs_interaction_to_db(self, client, test_db_pool):
diff --git a/python/tests/unit/db/test_repository.py b/python/tests/unit/db/test_repository.py
diff --git a/python/uv.lock b/python/uv.lock
