Skip to content

Commit f1613a3

Browse files
Trojan-Dev-AFKTrojan-Dev-AFK
committed
chat persistance
1 parent 2dfd1da commit f1613a3

26 files changed

Lines changed: 1420 additions & 268 deletions

‎README.md‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
# Agentic RAG API
22

33
A multi-tenant REST API that lets companies upload PDF documents and interrogate them
4-
through a conversational AI agent. Each company's data is fully isolated.
4+
through a conversational AI agent. Each company's data is fully isolated, and chat
5+
conversation history is persisted per user.
56

67
---
78

@@ -57,6 +58,8 @@ Full setup guide: [docs/RUNBOOK.md](docs/RUNBOOK.md)
5758
| `GET` | `/v1/documents/{id}` | Admin (own company) or Super Admin |
5859
| `DELETE` | `/v1/documents/{id}` | Admin (own company) or Super Admin |
5960
| `POST` | `/v1/chat/invoke` | Admin or Employee |
61+
| `GET` | `/v1/chat/conversations` | Admin or Employee |
62+
| `GET` | `/v1/chat/conversations/{conversation_id}/messages` | Admin or Employee |
6063

6164
---
6265

@@ -70,6 +73,8 @@ Full setup guide: [docs/RUNBOOK.md](docs/RUNBOOK.md)
7073
| Migrations | Alembic |
7174
| Database | PostgreSQL 16 + pgvector |
7275
| Vector search | pgvector cosine distance |
76+
| PDF extraction | pypdf + pdfplumber (table-aware extraction) |
77+
| OCR fallback | rapidocr-onnxruntime + pypdfium2 (for image/scanned pages) |
7378
| Embeddings | `sentence-transformers` (`all-MiniLM-L6-v2`, 384-dim) |
7479
| LLM | Ollama (`llama3.1`) via `langchain-ollama` |
7580
| Agent framework | LangGraph |
@@ -128,9 +133,12 @@ Full variable reference: [docs/RUNBOOK.md#environment-variables-reference](docs/
128133

129134
| Document | Contents |
130135
|----------|----------|
136+
| [docs/API_DOCS.md](docs/API_DOCS.md) | Endpoint-level request/response reference with role-based access notes |
131137
| [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) | Component diagram, request lifecycle, singleton model loading, storage backend design, RBAC model |
138+
| [docs/ENDPOINT_ACCESS_MATRIX.md](docs/ENDPOINT_ACCESS_MATRIX.md) | Quick role/access lookup for all endpoints |
132139
| [docs/RUNBOOK.md](docs/RUNBOOK.md) | Full setup, service startup/shutdown, migrations, monitoring, common ops tasks, S3 setup, backup/restore |
133140
| [docs/DEVELOPMENT.md](docs/DEVELOPMENT.md) | Code standards, docstring style, singleton pattern, adding endpoints, migrations, logging conventions |
141+
| [docs/CHANGELOG.md](docs/CHANGELOG.md) | Historical record of architecture, ingestion, RBAC, CI, and chat-history milestones |
134142

135143
---
136144

‎alembic/versions/0003_add_document_chunk_type.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
"""add document chunk type
2+
3+
Revision ID: 0003
4+
Revises: 0002
5+
Create Date: 2026-06-27
6+
"""
7+
8+
from collections.abc import Sequence
9+
from typing import Union
10+
11+
import sqlalchemy as sa
12+
13+
from alembic import op
14+
15+
revision: str = "0003"
16+
down_revision: Union[str, None] = "0002"
17+
branch_labels: Union[str, Sequence[str], None] = None
18+
depends_on: Union[str, Sequence[str], None] = None
19+
20+
21+
def upgrade() -> None:
22+
op.execute(sa.text("""
23+
DO $$ BEGIN
24+
IF EXISTS (SELECT 1 FROM pg_type WHERE typname = 'chunktype') THEN
25+
IF NOT EXISTS (
26+
SELECT 1 FROM pg_enum e
27+
JOIN pg_type t ON t.oid = e.enumtypid
28+
WHERE t.typname = 'chunktype' AND e.enumlabel = 'TABLE'
29+
) THEN
30+
DROP TYPE chunktype CASCADE;
31+
CREATE TYPE chunktype AS ENUM ('TEXT', 'TABLE');
32+
END IF;
33+
ELSE
34+
CREATE TYPE chunktype AS ENUM ('TEXT', 'TABLE');
35+
END IF;
36+
END $$
37+
"""))
38+
39+
op.execute(sa.text("""
40+
ALTER TABLE document_chunks
41+
ADD COLUMN IF NOT EXISTS chunk_type chunktype NOT NULL DEFAULT 'TEXT'
42+
"""))
43+
44+
45+
def downgrade() -> None:
46+
op.execute(sa.text("""
47+
ALTER TABLE document_chunks
48+
DROP COLUMN IF EXISTS chunk_type
49+
"""))
50+
op.execute(sa.text("DROP TYPE IF EXISTS chunktype"))

‎alembic/versions/0004_add_ocr_chunk_type.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
"""add ocr chunk type
2+
3+
Revision ID: 0004
4+
Revises: 0003
5+
Create Date: 2026-06-27
6+
"""
7+
8+
from collections.abc import Sequence
9+
from typing import Union
10+
11+
import sqlalchemy as sa
12+
13+
from alembic import op
14+
15+
revision: str = "0004"
16+
down_revision: Union[str, None] = "0003"
17+
branch_labels: Union[str, Sequence[str], None] = None
18+
depends_on: Union[str, Sequence[str], None] = None
19+
20+
21+
def upgrade() -> None:
22+
op.execute(sa.text("""
23+
DO $$ BEGIN
24+
IF NOT EXISTS (
25+
SELECT 1 FROM pg_enum e
26+
JOIN pg_type t ON t.oid = e.enumtypid
27+
WHERE t.typname = 'chunktype' AND e.enumlabel = 'OCR'
28+
) THEN
29+
ALTER TYPE chunktype ADD VALUE 'OCR';
30+
END IF;
31+
END $$
32+
"""))
33+
34+
35+
def downgrade() -> None:
36+
# PostgreSQL does not support dropping enum values directly.
37+
# Downgrade is intentionally a no-op for enum label removal.
38+
pass

‎alembic/versions/0005_add_chat_history_tables.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
"""add chat history tables
2+
3+
Revision ID: 0005
4+
Revises: 0004
5+
Create Date: 2026-06-27
6+
"""
7+
8+
from collections.abc import Sequence
9+
from typing import Union
10+
11+
import sqlalchemy as sa
12+
13+
from alembic import op
14+
15+
revision: str = "0005"
16+
down_revision: Union[str, None] = "0004"
17+
branch_labels: Union[str, Sequence[str], None] = None
18+
depends_on: Union[str, Sequence[str], None] = None
19+
20+
21+
def upgrade() -> None:
22+
op.execute(sa.text("""
23+
DO $$ BEGIN
24+
IF EXISTS (SELECT 1 FROM pg_type WHERE typname = 'chatmessagerole') THEN
25+
IF NOT EXISTS (
26+
SELECT 1 FROM pg_enum e
27+
JOIN pg_type t ON t.oid = e.enumtypid
28+
WHERE t.typname = 'chatmessagerole' AND e.enumlabel = 'assistant'
29+
) THEN
30+
DROP TYPE chatmessagerole CASCADE;
31+
CREATE TYPE chatmessagerole AS ENUM ('user', 'assistant');
32+
END IF;
33+
ELSE
34+
CREATE TYPE chatmessagerole AS ENUM ('user', 'assistant');
35+
END IF;
36+
END $$
37+
"""))
38+
39+
op.execute(sa.text("""
40+
CREATE TABLE IF NOT EXISTS chat_conversations (
41+
id VARCHAR NOT NULL,
42+
user_id VARCHAR NOT NULL REFERENCES users(id) ON DELETE CASCADE,
43+
company_id VARCHAR NOT NULL REFERENCES companies(id) ON DELETE CASCADE,
44+
title VARCHAR,
45+
created_at TIMESTAMPTZ DEFAULT now(),
46+
updated_at TIMESTAMPTZ DEFAULT now(),
47+
PRIMARY KEY (id)
48+
)
49+
"""))
50+
op.execute(sa.text("CREATE INDEX IF NOT EXISTS ix_chat_conversations_user_id ON chat_conversations (user_id)"))
51+
op.execute(
52+
sa.text("CREATE INDEX IF NOT EXISTS ix_chat_conversations_company_id ON chat_conversations (company_id)")
53+
)
54+
55+
op.execute(sa.text("""
56+
CREATE TABLE IF NOT EXISTS chat_messages (
57+
id VARCHAR NOT NULL,
58+
conversation_id VARCHAR NOT NULL REFERENCES chat_conversations(id) ON DELETE CASCADE,
59+
role chatmessagerole NOT NULL,
60+
content TEXT NOT NULL,
61+
created_at TIMESTAMPTZ DEFAULT now(),
62+
PRIMARY KEY (id)
63+
)
64+
"""))
65+
op.execute(
66+
sa.text("CREATE INDEX IF NOT EXISTS ix_chat_messages_conversation_id ON chat_messages (conversation_id)")
67+
)
68+
69+
70+
def downgrade() -> None:
71+
op.execute(sa.text("DROP TABLE IF EXISTS chat_messages"))
72+
op.execute(sa.text("DROP TABLE IF EXISTS chat_conversations"))
73+
op.execute(sa.text("DROP TYPE IF EXISTS chatmessagerole"))

‎alembic/versions/0006_chat_messages_single_turn_rows.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
"""chat messages single-turn rows
2+
3+
Revision ID: 0006
4+
Revises: 0005
5+
Create Date: 2026-06-27
6+
"""
7+
8+
from collections.abc import Sequence
9+
from typing import Union
10+
11+
import sqlalchemy as sa
12+
13+
from alembic import op
14+
15+
revision: str = "0006"
16+
down_revision: Union[str, None] = "0005"
17+
branch_labels: Union[str, Sequence[str], None] = None
18+
depends_on: Union[str, Sequence[str], None] = None
19+
20+
21+
def upgrade() -> None:
22+
op.execute(sa.text("""
23+
ALTER TABLE chat_messages
24+
ADD COLUMN IF NOT EXISTS user_query TEXT NOT NULL DEFAULT '',
25+
ADD COLUMN IF NOT EXISTS assistant_response TEXT NOT NULL DEFAULT ''
26+
"""))
27+
28+
# Backfill legacy two-row format into the new single-turn fields where possible.
29+
op.execute(sa.text("""
30+
UPDATE chat_messages
31+
SET user_query = content
32+
WHERE role::text = 'user' AND user_query = ''
33+
"""))
34+
op.execute(sa.text("""
35+
UPDATE chat_messages
36+
SET assistant_response = content
37+
WHERE role::text = 'assistant' AND assistant_response = ''
38+
"""))
39+
40+
# Keep legacy columns for backward compatibility, but allow new writes without them.
41+
op.execute(sa.text("""
42+
ALTER TABLE chat_messages
43+
ALTER COLUMN role DROP NOT NULL,
44+
ALTER COLUMN content DROP NOT NULL
45+
"""))
46+
47+
48+
def downgrade() -> None:
49+
op.execute(sa.text("""
50+
ALTER TABLE chat_messages
51+
ALTER COLUMN role SET NOT NULL,
52+
ALTER COLUMN content SET NOT NULL
53+
"""))
54+
op.execute(sa.text("ALTER TABLE chat_messages DROP COLUMN IF EXISTS assistant_response"))
55+
op.execute(sa.text("ALTER TABLE chat_messages DROP COLUMN IF EXISTS user_query"))

‎alembic/versions/0007_remove_legacy_chat_message_columns.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
"""remove legacy chat message columns
2+
3+
Revision ID: 0007
4+
Revises: 0006
5+
Create Date: 2026-06-27
6+
"""
7+
8+
from collections.abc import Sequence
9+
from typing import Union
10+
11+
import sqlalchemy as sa
12+
13+
from alembic import op
14+
15+
revision: str = "0007"
16+
down_revision: Union[str, None] = "0006"
17+
branch_labels: Union[str, Sequence[str], None] = None
18+
depends_on: Union[str, Sequence[str], None] = None
19+
20+
21+
def upgrade() -> None:
22+
op.execute(sa.text("ALTER TABLE chat_messages DROP COLUMN IF EXISTS role"))
23+
op.execute(sa.text("ALTER TABLE chat_messages DROP COLUMN IF EXISTS content"))
24+
op.execute(sa.text("DROP TYPE IF EXISTS chatmessagerole"))
25+
26+
27+
def downgrade() -> None:
28+
op.execute(sa.text("CREATE TYPE chatmessagerole AS ENUM ('user', 'assistant')"))
29+
op.execute(sa.text("""
30+
ALTER TABLE chat_messages
31+
ADD COLUMN IF NOT EXISTS role chatmessagerole,
32+
ADD COLUMN IF NOT EXISTS content TEXT
33+
"""))
34+
op.execute(sa.text("""
35+
UPDATE chat_messages
36+
SET role = 'user', content = user_query
37+
WHERE content IS NULL
38+
"""))

‎app/agent/tools/vector_search.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
from app.core.config import settings
1111
from app.core.logger import get_logger
12-
from app.db.models import Document, DocumentChunk
12+
from app.db.models import ChunkType, Document, DocumentChunk
1313
from app.db.session import AsyncSessionLocal
1414

1515
logger = get_logger(__name__)
@@ -103,14 +103,14 @@ async def search_documents(query: str) -> str:
103103
try:
104104
async with AsyncSessionLocal() as db:
105105
stmt = (
106-
select(DocumentChunk.text_content)
106+
select(DocumentChunk.text_content, DocumentChunk.chunk_type)
107107
.join(Document, Document.id == DocumentChunk.document_id)
108108
.where(Document.company_id == company_id)
109109
.order_by(DocumentChunk.embedding.cosine_distance(query_vector))
110110
.limit(5)
111111
)
112112
result = await db.execute(stmt)
113-
chunks = result.scalars().all()
113+
chunks = result.all()
114114
except SQLAlchemyError as exc:
115115
logger.error("Vector search database query failed", exc_info=exc)
116116
return "Search is temporarily unavailable due to a database error."
@@ -120,4 +120,12 @@ async def search_documents(query: str) -> str:
120120
return "No relevant information found in the documents."
121121

122122
logger.info("Vector search completed", extra={"results": len(chunks), "company_id": company_id})
123-
return "\n\n---\n\n".join(chunks)
123+
formatted_chunks = [
124+
(
125+
f"[TABLE]\n{text}"
126+
if chunk_type == ChunkType.TABLE
127+
else f"[OCR]\n{text}" if chunk_type == ChunkType.OCR else text
128+
)
129+
for text, chunk_type in chunks
130+
]
131+
return "\n\n---\n\n".join(formatted_chunks)

‎app/api/v1/endpoints/chat.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
"""Chat endpoint — submits a query to the LangGraph RAG agent."""
22

33
from fastapi import APIRouter, Depends
4+
from sqlalchemy.ext.asyncio import AsyncSession
45

56
from app.api.dependencies import require_company_user
67
from app.db.models import User
7-
from app.schemas.chat import ChatRequest, ChatResponse
8+
from app.db.session import get_db
9+
from app.schemas.chat import ChatConversationResponse, ChatMessageResponse, ChatRequest, ChatResponse
810
from app.services import chat_service
911

1012
router = APIRouter()
@@ -20,6 +22,7 @@
2022
)
2123
async def invoke_agent(
2224
request: ChatRequest,
25+
db: AsyncSession = Depends(get_db),
2326
current_user: User = Depends(require_company_user),
2427
):
2528
"""
@@ -31,5 +34,33 @@ async def invoke_agent(
3134
3235
**Access:** `admin` and `employee` only. `super_admin` receives **403**.
3336
"""
34-
response_text = await chat_service.invoke_agent(query=request.query, current_user=current_user)
35-
return ChatResponse(response=response_text)
37+
response_text, conversation_id = await chat_service.invoke_agent(
38+
query=request.query,
39+
conversation_id=request.conversation_id,
40+
db=db,
41+
current_user=current_user,
42+
)
43+
return ChatResponse(response=response_text, conversation_id=conversation_id)
44+
45+
46+
@router.get("/conversations", response_model=list[ChatConversationResponse])
47+
async def list_conversations(
48+
db: AsyncSession = Depends(get_db),
49+
current_user: User = Depends(require_company_user),
50+
):
51+
"""List persisted chat conversations for the current authenticated user."""
52+
return await chat_service.list_conversations(db=db, current_user=current_user)
53+
54+
55+
@router.get("/conversations/{conversation_id}/messages", response_model=list[ChatMessageResponse])
56+
async def get_conversation_messages(
57+
conversation_id: str,
58+
db: AsyncSession = Depends(get_db),
59+
current_user: User = Depends(require_company_user),
60+
):
61+
"""Get persisted messages for one conversation owned by the current user."""
62+
return await chat_service.get_conversation_messages(
63+
conversation_id=conversation_id,
64+
db=db,
65+
current_user=current_user,
66+
)

0 commit comments

Comments
 (0)