fix: [hotfix] return schemas, not models

This commit is contained in:
Benjamin McCormick 2025-10-09 17:25:11 -04:00
parent a657f88951
commit 1eefc4af31
2 changed files with 67 additions and 13 deletions

View File

@ -7,7 +7,7 @@ from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import apaginate
from sqlalchemy.ext.asyncio import AsyncSession
from src import config, crud, models, schemas
from src import config, crud, schemas
from src.dependencies import db, tracked_db
from src.exceptions import (
AuthenticationException,
@ -87,7 +87,7 @@ async def _get_session_context_task(
session_id: str,
token_limit: int,
include_summary: bool,
) -> tuple[schemas.Summary | None, list[models.Message]]:
) -> tuple[schemas.Summary | None, list[schemas.Message]]:
"""
Atomic task to get session context using tracked_db.
@ -101,13 +101,16 @@ async def _get_session_context_task(
Tuple of (summary, messages)
"""
async with tracked_db("get_session_context") as db:
return await summarizer.get_session_context(
summary, messages = await summarizer.get_session_context(
db,
workspace_name=workspace_id,
session_name=session_id,
token_limit=token_limit,
include_summary=include_summary,
)
# Convert SQLAlchemy models to Pydantic schemas while session is active
message_schemas = [schemas.Message.model_validate(msg) for msg in messages]
return summary, message_schemas
@router.post(
@ -501,7 +504,7 @@ async def get_session_context(
)
return schemas.SessionContext(
name=session_id,
messages=messages, # pyright: ignore -- db message type and schema message type are different, but excess gets removed by schema
messages=messages,
summary=summary,
)
@ -540,7 +543,7 @@ async def get_session_context(
return schemas.SessionContext(
name=session_id,
messages=messages, # pyright: ignore -- db message type and schema message type are different, but excess gets removed by schema
messages=messages,
summary=summary,
peer_representation=representation,
peer_card=card,

View File

@ -44,6 +44,7 @@ Optional arguments:
--json-output: Path to write JSON summary results for analytics (if not provided, creates timestamped file in tests/bench/eval_results)
--merge-sessions: Merge all sessions within a question into a single session (default: False)
--cleanup-workspace: Delete workspace after executing each question (default: False)
--use-get-context: Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)
```
## Other notes
@ -59,10 +60,11 @@ import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from typing import Any, cast
import tiktoken
from anthropic import AsyncAnthropic
from anthropic.types import MessageParam
from dotenv import load_dotenv
from honcho import AsyncHoncho
from honcho.async_client.session import SessionPeerConfig
@ -123,6 +125,7 @@ class LongMemEvalRunner:
timeout_seconds: int | None = None,
merge_sessions: bool = False,
cleanup_workspace: bool = False,
use_get_context: bool = False,
):
"""
Initialize the test runner.
@ -134,6 +137,7 @@ class LongMemEvalRunner:
timeout_seconds: Timeout for deriver queue in seconds
merge_sessions: If True, merge all sessions within a question into one session
cleanup_workspace: If True, delete workspace after executing question (default: False)
use_get_context: If True, use get_context + judge LLM instead of dialectic .chat endpoint
"""
self.base_api_port: int = base_api_port
self.pool_size: int = pool_size
@ -143,6 +147,7 @@ class LongMemEvalRunner:
)
self.merge_sessions: bool = merge_sessions
self.cleanup_workspace: bool = cleanup_workspace
self.use_get_context: bool = use_get_context
# Initialize metrics collector
self.metrics_collector: MetricsCollector = MetricsCollector()
@ -396,7 +401,7 @@ Actual response: "{actual_response}"
Evaluate whether the actual response correctly answers the question based on the expected answer. Focus on factual accuracy and evidence that the AI accessed the correct memory."""
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-20250514",
model="claude-sonnet-4-5",
max_tokens=300,
temperature=0.0,
system=system_prompt,
@ -621,6 +626,7 @@ Evaluate whether the actual response correctly answers the question based on the
)
)
else:
merged_session_id = None
# create separate sessions
# Zip together dates, session IDs, and session content
for session_date, session_id, session_messages in zip(
@ -732,13 +738,51 @@ Evaluate whether the actual response correctly answers the question based on the
output_lines.append(f"\nAsking question: {question_with_date}")
try:
# Use the appropriate peer based on question type
if is_assistant_type:
# For assistant questions, use the assistant peer
actual_response = await assistant_peer.chat(question_with_date)
if self.use_get_context:
# Use get_context instead of dialectic .chat endpoint
# Get the session to retrieve context from
if not self.merge_sessions or merged_session_id is None:
raise ValueError(
"Merged session ID is required when using get_context. Set --merge-sessions to True."
)
session = await honcho_client.session(id=merged_session_id)
# Get context for the appropriate peer
peer_id = "assistant" if is_assistant_type else "user"
context = await session.get_context(
summary=True,
peer_target=peer_id,
last_user_message=question,
)
# Format context using to_anthropic method
context_messages = context.to_anthropic(assistant="assistant")
# Add the question as the final user message
context_messages.append(
{"role": "user", "content": question_with_date}
)
# Call Anthropic API to generate response
response = await self.anthropic_client.messages.create(
model="claude-sonnet-4-5",
max_tokens=1024,
messages=cast(list[MessageParam], context_messages),
)
if not response.content:
raise ValueError("Anthropic returned empty response")
content_block = response.content[0]
actual_response = getattr(content_block, "text", "")
else:
# For user questions, use the user peer (default behavior)
actual_response = await user_peer.chat(question_with_date)
# Use the appropriate peer based on question type
if is_assistant_type:
# For assistant questions, use the assistant peer
actual_response = await assistant_peer.chat(question_with_date)
else:
# For user questions, use the user peer (default behavior)
actual_response = await user_peer.chat(question_with_date)
# Clean up workspace if requested
if self.cleanup_workspace:
@ -1158,6 +1202,12 @@ Examples:
help="Delete workspace after executing each question (default: False)",
)
parser.add_argument(
"--use-get-context",
action="store_true",
help="Use get_context + judge LLM instead of dialectic .chat endpoint (default: False)",
)
args = parser.parse_args()
# Validate arguments
@ -1181,6 +1231,7 @@ Examples:
timeout_seconds=args.timeout,
merge_sessions=args.merge_sessions,
cleanup_workspace=args.cleanup_workspace,
use_get_context=args.use_get_context,
)
try: