Dynamic session context to be used in the conversational agent
This commit is contained in:
parent
6f94f71d0c
commit
aaddf4a5a9
11 changed files with 782 additions and 5 deletions
|
|
@ -17,6 +17,7 @@ from haiku.rag.agents.chat import (
|
||||||
ChatDeps,
|
ChatDeps,
|
||||||
ChatSessionState,
|
ChatSessionState,
|
||||||
QAResponse,
|
QAResponse,
|
||||||
|
SessionContext,
|
||||||
create_chat_agent,
|
create_chat_agent,
|
||||||
)
|
)
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
|
|
@ -83,6 +84,7 @@ async def stream_chat(request: Request) -> Response:
|
||||||
# Restore session state from incoming AG-UI state (look under namespaced key)
|
# Restore session state from incoming AG-UI state (look under namespaced key)
|
||||||
initial_qa_history: list[QAResponse] = []
|
initial_qa_history: list[QAResponse] = []
|
||||||
background_context: str | None = None
|
background_context: str | None = None
|
||||||
|
session_context: SessionContext | None = None
|
||||||
state = getattr(run_input, "state", None)
|
state = getattr(run_input, "state", None)
|
||||||
if state:
|
if state:
|
||||||
chat_state = state.get(AGUI_STATE_KEY, state)
|
chat_state = state.get(AGUI_STATE_KEY, state)
|
||||||
|
|
@ -91,6 +93,9 @@ async def stream_chat(request: Request) -> Response:
|
||||||
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
|
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
|
||||||
]
|
]
|
||||||
background_context = chat_state.get("background_context")
|
background_context = chat_state.get("background_context")
|
||||||
|
ctx_data = chat_state.get("session_context")
|
||||||
|
if ctx_data and isinstance(ctx_data, dict):
|
||||||
|
session_context = SessionContext(**ctx_data)
|
||||||
|
|
||||||
# Build deps with session state
|
# Build deps with session state
|
||||||
thread_id = getattr(run_input, "thread_id", None)
|
thread_id = getattr(run_input, "thread_id", None)
|
||||||
|
|
@ -101,6 +106,7 @@ async def stream_chat(request: Request) -> Response:
|
||||||
session_id=thread_id or "",
|
session_id=thread_id or "",
|
||||||
qa_history=initial_qa_history,
|
qa_history=initial_qa_history,
|
||||||
background_context=background_context,
|
background_context=background_context,
|
||||||
|
session_context=session_context,
|
||||||
),
|
),
|
||||||
state_key=AGUI_STATE_KEY,
|
state_key=AGUI_STATE_KEY,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
from haiku.rag.agents.chat.agent import create_chat_agent
|
from haiku.rag.agents.chat.agent import create_chat_agent
|
||||||
|
from haiku.rag.agents.chat.context import summarize_session, update_session_context
|
||||||
from haiku.rag.agents.chat.search import SearchAgent
|
from haiku.rag.agents.chat.search import SearchAgent
|
||||||
from haiku.rag.agents.chat.state import (
|
from haiku.rag.agents.chat.state import (
|
||||||
AGUI_STATE_KEY,
|
AGUI_STATE_KEY,
|
||||||
|
|
@ -7,6 +8,7 @@ from haiku.rag.agents.chat.state import (
|
||||||
CitationInfo,
|
CitationInfo,
|
||||||
QAResponse,
|
QAResponse,
|
||||||
SearchDeps,
|
SearchDeps,
|
||||||
|
SessionContext,
|
||||||
build_document_filter,
|
build_document_filter,
|
||||||
format_conversation_context,
|
format_conversation_context,
|
||||||
)
|
)
|
||||||
|
|
@ -20,6 +22,9 @@ __all__ = [
|
||||||
"CitationInfo",
|
"CitationInfo",
|
||||||
"QAResponse",
|
"QAResponse",
|
||||||
"SearchDeps",
|
"SearchDeps",
|
||||||
|
"SessionContext",
|
||||||
"build_document_filter",
|
"build_document_filter",
|
||||||
"format_conversation_context",
|
"format_conversation_context",
|
||||||
|
"summarize_session",
|
||||||
|
"update_session_context",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,10 @@
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
|
||||||
from ag_ui.core import EventType, StateSnapshotEvent
|
from ag_ui.core import EventType, StateSnapshotEvent
|
||||||
from pydantic_ai import Agent, RunContext, ToolReturn
|
from pydantic_ai import Agent, RunContext, ToolReturn
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.context import update_session_context
|
||||||
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
|
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
|
||||||
from haiku.rag.agents.chat.search import SearchAgent
|
from haiku.rag.agents.chat.search import SearchAgent
|
||||||
from haiku.rag.agents.chat.state import (
|
from haiku.rag.agents.chat.state import (
|
||||||
|
|
@ -19,6 +23,25 @@ from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||||
from haiku.rag.config.models import AppConfig
|
from haiku.rag.config.models import AppConfig
|
||||||
from haiku.rag.utils import get_model
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
async def _update_context_background(
|
||||||
|
qa_history: list[QAResponse],
|
||||||
|
config: AppConfig,
|
||||||
|
session_state: ChatSessionState,
|
||||||
|
) -> None:
|
||||||
|
"""Background task to update session context after an ask."""
|
||||||
|
try:
|
||||||
|
await update_session_context(
|
||||||
|
qa_history=qa_history,
|
||||||
|
config=config,
|
||||||
|
session_state=session_state,
|
||||||
|
)
|
||||||
|
logger.debug("Session context updated successfully")
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Failed to update session context")
|
||||||
|
|
||||||
|
|
||||||
def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
"""Create the chat agent with search and ask tools."""
|
"""Create the chat agent with search and ask tools."""
|
||||||
|
|
@ -98,6 +121,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
if ctx.deps.session_state
|
if ctx.deps.session_state
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
|
session_context=(
|
||||||
|
ctx.deps.session_state.session_context
|
||||||
|
if ctx.deps.session_state
|
||||||
|
else None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Return detailed results for the agent to present
|
# Return detailed results for the agent to present
|
||||||
|
|
@ -194,11 +222,20 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
# Build and run the conversational research graph
|
# Build and run the conversational research graph
|
||||||
graph = build_conversational_graph(config=ctx.deps.config)
|
graph = build_conversational_graph(config=ctx.deps.config)
|
||||||
|
|
||||||
background_context = (
|
# Determine background context:
|
||||||
ctx.deps.session_state.background_context
|
# 1. Use session_context summary if available (compressed history)
|
||||||
if ctx.deps.session_state
|
# 2. Fall back to explicit background_context if set
|
||||||
else None
|
background_context: str | None = None
|
||||||
)
|
if ctx.deps.session_state:
|
||||||
|
if (
|
||||||
|
ctx.deps.session_state.session_context
|
||||||
|
and ctx.deps.session_state.session_context.summary
|
||||||
|
):
|
||||||
|
background_context = (
|
||||||
|
ctx.deps.session_state.session_context.render_markdown()
|
||||||
|
)
|
||||||
|
elif ctx.deps.session_state.background_context:
|
||||||
|
background_context = ctx.deps.session_state.background_context
|
||||||
|
|
||||||
context = ResearchContext(
|
context = ResearchContext(
|
||||||
original_question=question,
|
original_question=question,
|
||||||
|
|
@ -248,6 +285,15 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
-MAX_QA_HISTORY:
|
-MAX_QA_HISTORY:
|
||||||
]
|
]
|
||||||
|
|
||||||
|
# Spawn background task to update session context
|
||||||
|
asyncio.create_task(
|
||||||
|
_update_context_background(
|
||||||
|
qa_history=list(ctx.deps.session_state.qa_history),
|
||||||
|
config=ctx.deps.config,
|
||||||
|
session_state=ctx.deps.session_state,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
# Build new state with citations AND accumulated qa_history
|
# Build new state with citations AND accumulated qa_history
|
||||||
new_state = ChatSessionState(
|
new_state = ChatSessionState(
|
||||||
session_id=(
|
session_id=(
|
||||||
|
|
@ -262,6 +308,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
if ctx.deps.session_state
|
if ctx.deps.session_state
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
|
session_context=(
|
||||||
|
ctx.deps.session_state.session_context
|
||||||
|
if ctx.deps.session_state
|
||||||
|
else None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Format answer with citation references and confidence
|
# Format answer with citation references and confidence
|
||||||
|
|
|
||||||
71
haiku_rag_slim/haiku/rag/agents/chat/context.py
Normal file
71
haiku_rag_slim/haiku/rag/agents/chat/context.py
Normal file
|
|
@ -0,0 +1,71 @@
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from pydantic_ai import Agent
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.prompts import SESSION_SUMMARY_PROMPT
|
||||||
|
from haiku.rag.agents.chat.state import ChatSessionState, QAResponse, SessionContext
|
||||||
|
from haiku.rag.config.models import AppConfig
|
||||||
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
|
|
||||||
|
async def summarize_session(
|
||||||
|
qa_history: list[QAResponse],
|
||||||
|
config: AppConfig,
|
||||||
|
) -> str:
|
||||||
|
"""Summarize qa_history into compact context.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
qa_history: List of Q&A pairs from the conversation.
|
||||||
|
config: AppConfig for model selection.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Markdown summary of the conversation history.
|
||||||
|
"""
|
||||||
|
if not qa_history:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
model = get_model(config.qa.model, config)
|
||||||
|
agent: Agent[None, str] = Agent(
|
||||||
|
model,
|
||||||
|
output_type=str,
|
||||||
|
instructions=SESSION_SUMMARY_PROMPT,
|
||||||
|
retries=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
history_text = _format_qa_history(qa_history)
|
||||||
|
result = await agent.run(history_text)
|
||||||
|
return result.output
|
||||||
|
|
||||||
|
|
||||||
|
async def update_session_context(
|
||||||
|
qa_history: list[QAResponse],
|
||||||
|
config: AppConfig,
|
||||||
|
session_state: ChatSessionState,
|
||||||
|
) -> None:
|
||||||
|
"""Update session context in the session state.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
qa_history: List of Q&A pairs from the conversation.
|
||||||
|
config: AppConfig for model selection.
|
||||||
|
session_state: The session state to update.
|
||||||
|
"""
|
||||||
|
summary = await summarize_session(qa_history, config)
|
||||||
|
session_state.session_context = SessionContext(
|
||||||
|
summary=summary,
|
||||||
|
last_updated=datetime.now(),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _format_qa_history(qa_history: list[QAResponse]) -> str:
|
||||||
|
"""Format qa_history for input to summarization."""
|
||||||
|
lines: list[str] = []
|
||||||
|
for i, qa in enumerate(qa_history, 1):
|
||||||
|
lines.append(f"## Q{i}: {qa.question}")
|
||||||
|
lines.append(f"**Answer** (confidence: {qa.confidence:.0%}):")
|
||||||
|
lines.append(qa.answer)
|
||||||
|
|
||||||
|
if qa.sources:
|
||||||
|
lines.append(f"**Sources:** {', '.join(qa.sources)}")
|
||||||
|
lines.append("")
|
||||||
|
|
||||||
|
return "\n".join(lines)
|
||||||
|
|
@ -38,3 +38,20 @@ For each user request:
|
||||||
You can optionally specify a limit parameter (default 5).
|
You can optionally specify a limit parameter (default 5).
|
||||||
|
|
||||||
IMPORTANT: You must make actual tool calls. Do not output "run_search(...)" as text."""
|
IMPORTANT: You must make actual tool calls. Do not output "run_search(...)" as text."""
|
||||||
|
|
||||||
|
SESSION_SUMMARY_PROMPT = """You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
|
||||||
|
|
||||||
|
Your summary should be concise (aim for 500-1500 tokens) and include:
|
||||||
|
|
||||||
|
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
|
||||||
|
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
|
||||||
|
3. **Current Focus** - What topic or question thread the user is currently exploring
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- Extract only high-signal information that would help answer follow-up questions
|
||||||
|
- Omit small talk, greetings, or low-confidence answers
|
||||||
|
- Use bullet points for clarity
|
||||||
|
- Keep technical details but compress verbose explanations
|
||||||
|
- Preserve document names/titles when mentioned in sources
|
||||||
|
|
||||||
|
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself."""
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
import hashlib
|
import hashlib
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
@ -55,6 +56,19 @@ class QAResponse(BaseModel):
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class SessionContext(BaseModel):
|
||||||
|
"""Compressed summary of conversation history for research graph."""
|
||||||
|
|
||||||
|
summary: str = ""
|
||||||
|
last_updated: datetime | None = None
|
||||||
|
|
||||||
|
def render_markdown(self) -> str:
|
||||||
|
"""Render context for injection into research graph."""
|
||||||
|
if not self.summary:
|
||||||
|
return ""
|
||||||
|
return f"## Prior Conversation Context\n\n{self.summary}"
|
||||||
|
|
||||||
|
|
||||||
class ChatSessionState(BaseModel):
|
class ChatSessionState(BaseModel):
|
||||||
"""State shared between frontend and agent via AG-UI."""
|
"""State shared between frontend and agent via AG-UI."""
|
||||||
|
|
||||||
|
|
@ -62,6 +76,7 @@ class ChatSessionState(BaseModel):
|
||||||
citations: list[CitationInfo] = []
|
citations: list[CitationInfo] = []
|
||||||
qa_history: list[QAResponse] = []
|
qa_history: list[QAResponse] = []
|
||||||
background_context: str | None = None
|
background_context: str | None = None
|
||||||
|
session_context: SessionContext | None = None
|
||||||
|
|
||||||
|
|
||||||
def format_conversation_context(qa_history: list[QAResponse]) -> str:
|
def format_conversation_context(qa_history: list[QAResponse]) -> str:
|
||||||
|
|
@ -206,6 +221,14 @@ class ChatDeps:
|
||||||
)
|
)
|
||||||
if "session_id" in state_data:
|
if "session_id" in state_data:
|
||||||
self.session_state.session_id = state_data.get("session_id", "")
|
self.session_state.session_id = state_data.get("session_id", "")
|
||||||
|
if "session_context" in state_data:
|
||||||
|
ctx_data = state_data.get("session_context")
|
||||||
|
if ctx_data is None:
|
||||||
|
self.session_state.session_context = None
|
||||||
|
elif isinstance(ctx_data, dict):
|
||||||
|
self.session_state.session_context = SessionContext(**ctx_data)
|
||||||
|
else:
|
||||||
|
self.session_state.session_context = ctx_data
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
|
||||||
229
tests/agents/chat/test_context.py
Normal file
229
tests/agents/chat/test_context.py
Normal file
|
|
@ -0,0 +1,229 @@
|
||||||
|
from datetime import datetime
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.state import (
|
||||||
|
CitationInfo,
|
||||||
|
QAResponse,
|
||||||
|
SessionContext,
|
||||||
|
)
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def vcr_cassette_dir():
|
||||||
|
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_context")
|
||||||
|
|
||||||
|
|
||||||
|
class TestSessionContext:
|
||||||
|
"""Tests for SessionContext model."""
|
||||||
|
|
||||||
|
def test_session_context_creation_empty(self):
|
||||||
|
"""Test SessionContext can be created with defaults."""
|
||||||
|
ctx = SessionContext()
|
||||||
|
assert ctx.summary == ""
|
||||||
|
assert ctx.last_updated is None
|
||||||
|
|
||||||
|
def test_session_context_creation_with_values(self):
|
||||||
|
"""Test SessionContext can be created with provided values."""
|
||||||
|
now = datetime.now()
|
||||||
|
ctx = SessionContext(
|
||||||
|
summary="User discussed authentication patterns.",
|
||||||
|
last_updated=now,
|
||||||
|
)
|
||||||
|
assert ctx.summary == "User discussed authentication patterns."
|
||||||
|
assert ctx.last_updated == now
|
||||||
|
|
||||||
|
def test_render_markdown_empty(self):
|
||||||
|
"""Test render_markdown returns empty string when no summary."""
|
||||||
|
ctx = SessionContext()
|
||||||
|
assert ctx.render_markdown() == ""
|
||||||
|
|
||||||
|
def test_render_markdown_with_summary(self):
|
||||||
|
"""Test render_markdown returns formatted markdown."""
|
||||||
|
ctx = SessionContext(
|
||||||
|
summary="## Key Facts\n- Authentication uses JWT\n- Rate limit is 100/min"
|
||||||
|
)
|
||||||
|
result = ctx.render_markdown()
|
||||||
|
assert "## Prior Conversation Context" in result
|
||||||
|
assert "Authentication uses JWT" in result
|
||||||
|
assert "Rate limit is 100/min" in result
|
||||||
|
|
||||||
|
def test_session_context_serialization_roundtrip(self):
|
||||||
|
"""Test SessionContext serializes and deserializes correctly."""
|
||||||
|
now = datetime.now()
|
||||||
|
original = SessionContext(
|
||||||
|
summary="Test summary with facts.",
|
||||||
|
last_updated=now,
|
||||||
|
)
|
||||||
|
# Serialize to dict
|
||||||
|
data = original.model_dump()
|
||||||
|
# Deserialize back
|
||||||
|
restored = SessionContext(**data)
|
||||||
|
|
||||||
|
assert restored.summary == original.summary
|
||||||
|
assert restored.last_updated == original.last_updated
|
||||||
|
|
||||||
|
|
||||||
|
class TestSummarizeSession:
|
||||||
|
"""Tests for summarize_session function."""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_summarize_session_empty_history(self, allow_model_requests):
|
||||||
|
"""Test summarize_session with empty qa_history returns empty string."""
|
||||||
|
from haiku.rag.agents.chat.context import summarize_session
|
||||||
|
|
||||||
|
result = await summarize_session(qa_history=[], config=Config)
|
||||||
|
assert result == ""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_summarize_session_single_entry(
|
||||||
|
self, allow_model_requests, temp_db_path
|
||||||
|
):
|
||||||
|
"""Test summarize_session with a single qa entry."""
|
||||||
|
from haiku.rag.agents.chat.context import summarize_session
|
||||||
|
|
||||||
|
qa_history = [
|
||||||
|
QAResponse(
|
||||||
|
question="What is the authentication method?",
|
||||||
|
answer="The API uses JWT tokens for authentication.",
|
||||||
|
confidence=0.95,
|
||||||
|
citations=[
|
||||||
|
CitationInfo(
|
||||||
|
index=1,
|
||||||
|
document_id="doc-1",
|
||||||
|
chunk_id="chunk-1",
|
||||||
|
document_uri="auth-guide.md",
|
||||||
|
document_title="Auth Guide",
|
||||||
|
content="JWT token details...",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
result = await summarize_session(qa_history=qa_history, config=Config)
|
||||||
|
|
||||||
|
# Should produce a non-empty summary
|
||||||
|
assert len(result) > 0
|
||||||
|
# Summary should mention authentication or JWT
|
||||||
|
assert "authentication" in result.lower() or "jwt" in result.lower()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_summarize_session_multiple_entries(
|
||||||
|
self, allow_model_requests, temp_db_path
|
||||||
|
):
|
||||||
|
"""Test summarize_session with multiple qa entries produces consolidated summary."""
|
||||||
|
from haiku.rag.agents.chat.context import summarize_session
|
||||||
|
|
||||||
|
qa_history = [
|
||||||
|
QAResponse(
|
||||||
|
question="What is the authentication method?",
|
||||||
|
answer="The API uses JWT tokens for authentication.",
|
||||||
|
confidence=0.95,
|
||||||
|
citations=[
|
||||||
|
CitationInfo(
|
||||||
|
index=1,
|
||||||
|
document_id="doc-1",
|
||||||
|
chunk_id="chunk-1",
|
||||||
|
document_uri="auth-guide.md",
|
||||||
|
document_title="Auth Guide",
|
||||||
|
content="JWT token details...",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
QAResponse(
|
||||||
|
question="What is the rate limit?",
|
||||||
|
answer="Rate limiting is set to 100 requests per minute.",
|
||||||
|
confidence=0.9,
|
||||||
|
citations=[
|
||||||
|
CitationInfo(
|
||||||
|
index=1,
|
||||||
|
document_id="doc-2",
|
||||||
|
chunk_id="chunk-2",
|
||||||
|
document_uri="api-reference.md",
|
||||||
|
document_title="API Reference",
|
||||||
|
content="Rate limit config...",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
QAResponse(
|
||||||
|
question="How do I refresh tokens?",
|
||||||
|
answer="Use the /refresh endpoint with your refresh token.",
|
||||||
|
confidence=0.85,
|
||||||
|
citations=[
|
||||||
|
CitationInfo(
|
||||||
|
index=1,
|
||||||
|
document_id="doc-1",
|
||||||
|
chunk_id="chunk-3",
|
||||||
|
document_uri="auth-guide.md",
|
||||||
|
document_title="Auth Guide",
|
||||||
|
content="Token refresh...",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
result = await summarize_session(qa_history=qa_history, config=Config)
|
||||||
|
|
||||||
|
# Should produce a non-empty summary
|
||||||
|
assert len(result) > 0
|
||||||
|
# Summary should contain structured sections
|
||||||
|
result_lower = result.lower()
|
||||||
|
assert "key facts" in result_lower or "established" in result_lower
|
||||||
|
assert "documents" in result_lower or "sources" in result_lower
|
||||||
|
|
||||||
|
|
||||||
|
class TestUpdateSessionContext:
|
||||||
|
"""Tests for update_session_context function."""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_update_session_context_updates_state(
|
||||||
|
self, allow_model_requests, temp_db_path
|
||||||
|
):
|
||||||
|
"""Test update_session_context updates the session_state."""
|
||||||
|
from haiku.rag.agents.chat.context import update_session_context
|
||||||
|
from haiku.rag.agents.chat.state import ChatSessionState
|
||||||
|
|
||||||
|
session_state = ChatSessionState(session_id="test-session")
|
||||||
|
|
||||||
|
qa_history = [
|
||||||
|
QAResponse(
|
||||||
|
question="What is the authentication method?",
|
||||||
|
answer="The API uses JWT tokens.",
|
||||||
|
confidence=0.95,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
await update_session_context(
|
||||||
|
qa_history=qa_history,
|
||||||
|
config=Config,
|
||||||
|
session_state=session_state,
|
||||||
|
)
|
||||||
|
|
||||||
|
# session_context should now be populated
|
||||||
|
assert session_state.session_context is not None
|
||||||
|
assert session_state.session_context.summary != ""
|
||||||
|
assert session_state.session_context.last_updated is not None
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_update_session_context_with_empty_history(self):
|
||||||
|
"""Test update_session_context with empty history sets empty context."""
|
||||||
|
from haiku.rag.agents.chat.context import update_session_context
|
||||||
|
from haiku.rag.agents.chat.state import ChatSessionState
|
||||||
|
|
||||||
|
session_state = ChatSessionState(session_id="test-session")
|
||||||
|
|
||||||
|
await update_session_context(
|
||||||
|
qa_history=[],
|
||||||
|
config=Config,
|
||||||
|
session_state=session_state,
|
||||||
|
)
|
||||||
|
|
||||||
|
# session_context should exist but have empty summary
|
||||||
|
assert session_state.session_context is not None
|
||||||
|
assert session_state.session_context.summary == ""
|
||||||
|
|
@ -470,3 +470,123 @@ def test_chat_deps_state_setter_with_citation_dicts():
|
||||||
assert citation.document_id == "doc-1"
|
assert citation.document_id == "doc-1"
|
||||||
assert citation.chunk_id == "chunk-1"
|
assert citation.chunk_id == "chunk-1"
|
||||||
assert citation.page_numbers == [1, 2]
|
assert citation.page_numbers == [1, 2]
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_deps_state_getter_includes_session_context():
|
||||||
|
"""Test ChatDeps.state getter includes session_context when present."""
|
||||||
|
from datetime import datetime
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.state import (
|
||||||
|
AGUI_STATE_KEY,
|
||||||
|
ChatDeps,
|
||||||
|
ChatSessionState,
|
||||||
|
SessionContext,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_config = MagicMock()
|
||||||
|
|
||||||
|
now = datetime.now()
|
||||||
|
session_state = ChatSessionState(
|
||||||
|
session_id="test-123",
|
||||||
|
session_context=SessionContext(
|
||||||
|
summary="User discussed authentication.",
|
||||||
|
last_updated=now,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
deps = ChatDeps(
|
||||||
|
client=mock_client,
|
||||||
|
config=mock_config,
|
||||||
|
session_state=session_state,
|
||||||
|
state_key=AGUI_STATE_KEY,
|
||||||
|
)
|
||||||
|
|
||||||
|
state = deps.state
|
||||||
|
assert state is not None
|
||||||
|
assert AGUI_STATE_KEY in state
|
||||||
|
assert state[AGUI_STATE_KEY]["session_context"] is not None
|
||||||
|
assert (
|
||||||
|
state[AGUI_STATE_KEY]["session_context"]["summary"]
|
||||||
|
== "User discussed authentication."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_deps_state_setter_restores_session_context():
|
||||||
|
"""Test ChatDeps.state setter restores session_context from incoming state."""
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
|
||||||
|
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_config = MagicMock()
|
||||||
|
|
||||||
|
session_state = ChatSessionState(session_id="initial")
|
||||||
|
deps = ChatDeps(
|
||||||
|
client=mock_client,
|
||||||
|
config=mock_config,
|
||||||
|
session_state=session_state,
|
||||||
|
state_key=AGUI_STATE_KEY,
|
||||||
|
)
|
||||||
|
|
||||||
|
incoming_state = {
|
||||||
|
AGUI_STATE_KEY: {
|
||||||
|
"session_id": "test-123",
|
||||||
|
"qa_history": [],
|
||||||
|
"citations": [],
|
||||||
|
"background_context": None,
|
||||||
|
"session_context": {
|
||||||
|
"summary": "Restored context summary.",
|
||||||
|
"last_updated": "2025-01-15T10:30:00",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
deps.state = incoming_state
|
||||||
|
|
||||||
|
assert deps.session_state is not None
|
||||||
|
assert deps.session_state.session_context is not None
|
||||||
|
assert deps.session_state.session_context.summary == "Restored context summary."
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_deps_state_setter_handles_null_session_context():
|
||||||
|
"""Test ChatDeps.state setter handles null session_context."""
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
from haiku.rag.agents.chat.state import (
|
||||||
|
AGUI_STATE_KEY,
|
||||||
|
ChatDeps,
|
||||||
|
ChatSessionState,
|
||||||
|
SessionContext,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_config = MagicMock()
|
||||||
|
|
||||||
|
# Start with a session_context
|
||||||
|
session_state = ChatSessionState(
|
||||||
|
session_id="test",
|
||||||
|
session_context=SessionContext(summary="Initial summary"),
|
||||||
|
)
|
||||||
|
deps = ChatDeps(
|
||||||
|
client=mock_client,
|
||||||
|
config=mock_config,
|
||||||
|
session_state=session_state,
|
||||||
|
state_key=AGUI_STATE_KEY,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Send null to clear it
|
||||||
|
incoming_state = {
|
||||||
|
AGUI_STATE_KEY: {
|
||||||
|
"session_id": "test",
|
||||||
|
"qa_history": [],
|
||||||
|
"citations": [],
|
||||||
|
"session_context": None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
deps.state = incoming_state
|
||||||
|
|
||||||
|
assert deps.session_state is not None
|
||||||
|
assert deps.session_state.session_context is None
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,94 @@
|
||||||
|
interactions:
|
||||||
|
- request:
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- application/json
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
content-length:
|
||||||
|
- '1493'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
host:
|
||||||
|
- localhost:11434
|
||||||
|
method: POST
|
||||||
|
parsed_body:
|
||||||
|
messages:
|
||||||
|
- content: |-
|
||||||
|
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
|
||||||
|
|
||||||
|
Your summary should be concise (aim for 500-1500 tokens) and include:
|
||||||
|
|
||||||
|
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
|
||||||
|
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
|
||||||
|
3. **Current Focus** - What topic or question thread the user is currently exploring
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- Extract only high-signal information that would help answer follow-up questions
|
||||||
|
- Omit small talk, greetings, or low-confidence answers
|
||||||
|
- Use bullet points for clarity
|
||||||
|
- Keep technical details but compress verbose explanations
|
||||||
|
- Preserve document names/titles when mentioned in sources
|
||||||
|
|
||||||
|
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
|
||||||
|
role: system
|
||||||
|
- content: |
|
||||||
|
## Q1: What is the authentication method?
|
||||||
|
**Answer** (confidence: 95%):
|
||||||
|
The API uses JWT tokens for authentication.
|
||||||
|
**Sources:** Auth Guide
|
||||||
|
|
||||||
|
## Q2: What is the rate limit?
|
||||||
|
**Answer** (confidence: 90%):
|
||||||
|
Rate limiting is set to 100 requests per minute.
|
||||||
|
**Sources:** API Reference
|
||||||
|
|
||||||
|
## Q3: How do I refresh tokens?
|
||||||
|
**Answer** (confidence: 85%):
|
||||||
|
Use the /refresh endpoint with your refresh token.
|
||||||
|
**Sources:** Auth Guide
|
||||||
|
role: user
|
||||||
|
model: gpt-oss
|
||||||
|
reasoning_effort: low
|
||||||
|
stream: false
|
||||||
|
uri: http://localhost:11434/v1/chat/completions
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '865'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
choices:
|
||||||
|
- finish_reason: stop
|
||||||
|
index: 0
|
||||||
|
message:
|
||||||
|
content: |-
|
||||||
|
## Key Facts Established
|
||||||
|
- **Authentication method**: JWT tokens.
|
||||||
|
- **Rate limit**: 100 requests per minute.
|
||||||
|
- **Token refresh**: use `/refresh` endpoint with a refresh token.
|
||||||
|
|
||||||
|
## Documents Referenced
|
||||||
|
- **Auth Guide** – contains details on JWT usage, token issuance, and refresh mechanism.
|
||||||
|
- **API Reference** – includes rate limiting policy and endpoint descriptions.
|
||||||
|
|
||||||
|
## Current Focus
|
||||||
|
The user is currently gathering foundational API usage details, specifically authentication methods, rate limits, and token refresh procedures.
|
||||||
|
reasoning: We need to summarize.
|
||||||
|
role: assistant
|
||||||
|
created: 1769007514
|
||||||
|
id: chatcmpl-760
|
||||||
|
model: gpt-oss
|
||||||
|
object: chat.completion
|
||||||
|
system_fingerprint: fp_ollama
|
||||||
|
usage:
|
||||||
|
completion_tokens: 121
|
||||||
|
prompt_tokens: 365
|
||||||
|
total_tokens: 486
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
version: 1
|
||||||
|
|
@ -0,0 +1,81 @@
|
||||||
|
interactions:
|
||||||
|
- request:
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- application/json
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
content-length:
|
||||||
|
- '1207'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
host:
|
||||||
|
- localhost:11434
|
||||||
|
method: POST
|
||||||
|
parsed_body:
|
||||||
|
messages:
|
||||||
|
- content: |-
|
||||||
|
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
|
||||||
|
|
||||||
|
Your summary should be concise (aim for 500-1500 tokens) and include:
|
||||||
|
|
||||||
|
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
|
||||||
|
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
|
||||||
|
3. **Current Focus** - What topic or question thread the user is currently exploring
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- Extract only high-signal information that would help answer follow-up questions
|
||||||
|
- Omit small talk, greetings, or low-confidence answers
|
||||||
|
- Use bullet points for clarity
|
||||||
|
- Keep technical details but compress verbose explanations
|
||||||
|
- Preserve document names/titles when mentioned in sources
|
||||||
|
|
||||||
|
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
|
||||||
|
role: system
|
||||||
|
- content: |
|
||||||
|
## Q1: What is the authentication method?
|
||||||
|
**Answer** (confidence: 95%):
|
||||||
|
The API uses JWT tokens for authentication.
|
||||||
|
**Sources:** Auth Guide
|
||||||
|
role: user
|
||||||
|
model: gpt-oss
|
||||||
|
reasoning_effort: low
|
||||||
|
stream: false
|
||||||
|
uri: http://localhost:11434/v1/chat/completions
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '575'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
choices:
|
||||||
|
- finish_reason: stop
|
||||||
|
index: 0
|
||||||
|
message:
|
||||||
|
content: |-
|
||||||
|
**Key Facts Established**
|
||||||
|
- The API uses **JWT tokens** for authentication.
|
||||||
|
|
||||||
|
**Documents Referenced**
|
||||||
|
- **Auth Guide** – Provides details on JWT usage for this API.
|
||||||
|
|
||||||
|
**Current Focus**
|
||||||
|
- Understanding the authentication method employed by the API.
|
||||||
|
reasoning: We need to summarize.
|
||||||
|
role: assistant
|
||||||
|
created: 1769007512
|
||||||
|
id: chatcmpl-739
|
||||||
|
model: gpt-oss
|
||||||
|
object: chat.completion
|
||||||
|
system_fingerprint: fp_ollama
|
||||||
|
usage:
|
||||||
|
completion_tokens: 65
|
||||||
|
prompt_tokens: 292
|
||||||
|
total_tokens: 357
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
version: 1
|
||||||
|
|
@ -0,0 +1,80 @@
|
||||||
|
interactions:
|
||||||
|
- request:
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- application/json
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
content-length:
|
||||||
|
- '1163'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
host:
|
||||||
|
- localhost:11434
|
||||||
|
method: POST
|
||||||
|
parsed_body:
|
||||||
|
messages:
|
||||||
|
- content: |-
|
||||||
|
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
|
||||||
|
|
||||||
|
Your summary should be concise (aim for 500-1500 tokens) and include:
|
||||||
|
|
||||||
|
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
|
||||||
|
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
|
||||||
|
3. **Current Focus** - What topic or question thread the user is currently exploring
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- Extract only high-signal information that would help answer follow-up questions
|
||||||
|
- Omit small talk, greetings, or low-confidence answers
|
||||||
|
- Use bullet points for clarity
|
||||||
|
- Keep technical details but compress verbose explanations
|
||||||
|
- Preserve document names/titles when mentioned in sources
|
||||||
|
|
||||||
|
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
|
||||||
|
role: system
|
||||||
|
- content: |
|
||||||
|
## Q1: What is the authentication method?
|
||||||
|
**Answer** (confidence: 95%):
|
||||||
|
The API uses JWT tokens.
|
||||||
|
role: user
|
||||||
|
model: gpt-oss
|
||||||
|
reasoning_effort: low
|
||||||
|
stream: false
|
||||||
|
uri: http://localhost:11434/v1/chat/completions
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '559'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
choices:
|
||||||
|
- finish_reason: stop
|
||||||
|
index: 0
|
||||||
|
message:
|
||||||
|
content: |-
|
||||||
|
## Key Facts Established
|
||||||
|
- The API authentication method is **JWT tokens** (high confidence 95%).
|
||||||
|
|
||||||
|
## Documents Referenced
|
||||||
|
- *None provided*.
|
||||||
|
|
||||||
|
## Current Focus
|
||||||
|
- The user is exploring details related to **API authentication mechanisms**.
|
||||||
|
reasoning: We need summary.
|
||||||
|
role: assistant
|
||||||
|
created: 1769007530
|
||||||
|
id: chatcmpl-975
|
||||||
|
model: gpt-oss
|
||||||
|
object: chat.completion
|
||||||
|
system_fingerprint: fp_ollama
|
||||||
|
usage:
|
||||||
|
completion_tokens: 64
|
||||||
|
prompt_tokens: 284
|
||||||
|
total_tokens: 348
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
version: 1
|
||||||
Loading…
Reference in a new issue