246 lines
7.5 KiB
Python
246 lines
7.5 KiB
Python
"""Request/response interceptor for OpenAI API calls."""
|
|
|
|
from typing import Any, Dict, List, Optional
|
|
from datetime import datetime
|
|
|
|
from hindsight_client import Hindsight
|
|
|
|
from .config import get_config, HindsightConfig
|
|
|
|
|
|
class HindsightInterceptor:
|
|
"""Intercepts OpenAI API calls to integrate with Hindsight."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the interceptor with a Hindsight client."""
|
|
self._client: Optional[Hindsight] = None
|
|
|
|
def get_client(self, config: HindsightConfig) -> Hindsight:
|
|
"""Get or create the Hindsight client."""
|
|
if self._client is None:
|
|
self._client = Hindsight(base_url=config.hindsight_api_url, timeout=30.0)
|
|
return self._client
|
|
|
|
def close(self):
|
|
"""Close the client."""
|
|
if self._client is not None:
|
|
self._client.close()
|
|
self._client = None
|
|
|
|
async def inject_memories(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
config: HindsightConfig,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Inject relevant memories into messages before sending to OpenAI.
|
|
"""
|
|
if not config.enabled:
|
|
return messages
|
|
|
|
try:
|
|
# Extract user query from messages
|
|
user_query = self._extract_user_query(messages)
|
|
if not user_query:
|
|
return messages
|
|
|
|
# Get client
|
|
client = self.get_client(config)
|
|
|
|
# Search for relevant memories
|
|
results = client.search(
|
|
agent_id=config.agent_id,
|
|
query=user_query,
|
|
max_tokens=4096,
|
|
thinking_budget=500,
|
|
)
|
|
|
|
if not results:
|
|
return messages
|
|
|
|
# Format memories and add to context
|
|
memory_context = self._format_memories(results)
|
|
|
|
# Add memory context to system message or create new one
|
|
updated_messages = self._add_memory_context(messages, memory_context)
|
|
|
|
return updated_messages
|
|
|
|
except Exception as e:
|
|
# Don't fail the request if memory retrieval fails
|
|
print(f"Warning: Failed to inject memories: {e}")
|
|
return messages
|
|
|
|
async def process_request(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
**kwargs: Any,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Process request before sending to OpenAI.
|
|
Retrieves relevant memories and adds them to the context.
|
|
"""
|
|
config = get_config()
|
|
return await self.inject_memories(messages, config)
|
|
|
|
async def store_conversation(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
response: Any,
|
|
config: HindsightConfig,
|
|
) -> None:
|
|
"""
|
|
Store conversation in Hindsight after receiving response from OpenAI.
|
|
"""
|
|
if not config.enabled or not config.store_conversations:
|
|
return
|
|
|
|
try:
|
|
# Extract conversation context
|
|
conversation = self._extract_conversation_context(messages, response)
|
|
if not conversation:
|
|
return
|
|
|
|
# Get client
|
|
client = self.get_client(config)
|
|
|
|
# Store conversation as memories
|
|
items = [
|
|
{
|
|
"content": msg["content"],
|
|
"context": f"role:{msg['role']}",
|
|
"event_date": datetime.now(),
|
|
}
|
|
for msg in conversation
|
|
]
|
|
|
|
# Store to Hindsight (async)
|
|
client.put_batch(
|
|
agent_id=config.agent_id,
|
|
items=items,
|
|
)
|
|
|
|
except Exception as e:
|
|
# Don't fail the request if storage fails
|
|
print(f"Warning: Failed to store conversation: {e}")
|
|
|
|
async def process_response(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
response: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""
|
|
Process response after receiving from OpenAI.
|
|
Stores the conversation in Hindsight.
|
|
"""
|
|
config = get_config()
|
|
await self.store_conversation(messages, response, config)
|
|
|
|
def _extract_user_query(self, messages: List[Dict[str, Any]]) -> Optional[str]:
|
|
"""Extract the user's query from messages."""
|
|
# Get the last user message
|
|
for msg in reversed(messages):
|
|
if msg.get("role") == "user":
|
|
content = msg.get("content")
|
|
if isinstance(content, str):
|
|
return content
|
|
elif isinstance(content, list):
|
|
# Handle structured content
|
|
for item in content:
|
|
if isinstance(item, dict) and item.get("type") == "text":
|
|
return item.get("text")
|
|
return None
|
|
|
|
def _format_memories(self, results: List[Dict[str, Any]]) -> str:
|
|
"""Format memory search results into context string."""
|
|
if not results:
|
|
return ""
|
|
|
|
memory_lines = []
|
|
for i, result in enumerate(results, 1):
|
|
text = result.get("text", "")
|
|
if text:
|
|
memory_lines.append(f"{i}. {text}")
|
|
|
|
if not memory_lines:
|
|
return ""
|
|
|
|
return (
|
|
"# Relevant Memories\n"
|
|
"The following memories may be relevant to this conversation:\n\n"
|
|
+ "\n".join(memory_lines)
|
|
)
|
|
|
|
def _add_memory_context(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
memory_context: str
|
|
) -> List[Dict[str, Any]]:
|
|
"""Add memory context to messages."""
|
|
if not memory_context:
|
|
return messages
|
|
|
|
# Check if there's already a system message
|
|
updated_messages = messages.copy()
|
|
|
|
for i, msg in enumerate(updated_messages):
|
|
if msg.get("role") == "system":
|
|
# Append to existing system message
|
|
updated_messages[i] = {
|
|
**msg,
|
|
"content": f"{msg['content']}\n\n{memory_context}"
|
|
}
|
|
return updated_messages
|
|
|
|
# No system message found, prepend one
|
|
system_message = {
|
|
"role": "system",
|
|
"content": memory_context
|
|
}
|
|
return [system_message] + updated_messages
|
|
|
|
def _extract_conversation_context(
|
|
self,
|
|
messages: List[Dict[str, Any]],
|
|
response: Any,
|
|
) -> List[Dict[str, str]]:
|
|
"""Extract conversation context for storage."""
|
|
conversation = []
|
|
|
|
# Add recent user messages
|
|
for msg in messages[-3:]: # Last 3 messages
|
|
if msg.get("role") in ("user", "assistant"):
|
|
content = msg.get("content")
|
|
if isinstance(content, str):
|
|
conversation.append({
|
|
"role": msg["role"],
|
|
"content": content,
|
|
})
|
|
|
|
# Add assistant response
|
|
if hasattr(response, "choices") and response.choices:
|
|
choice = response.choices[0]
|
|
if hasattr(choice, "message"):
|
|
content = choice.message.content
|
|
if content:
|
|
conversation.append({
|
|
"role": "assistant",
|
|
"content": content,
|
|
})
|
|
|
|
return conversation
|
|
|
|
|
|
# Global interceptor instance
|
|
_interceptor = HindsightInterceptor()
|
|
|
|
|
|
def get_interceptor() -> HindsightInterceptor:
|
|
"""Get the global interceptor instance."""
|
|
return _interceptor
|
|
|
|
|
|
def cleanup_interceptor() -> None:
|
|
"""Cleanup the global interceptor instance."""
|
|
_interceptor.close()
|