diff --git a/fly.toml b/fly.toml index 09c90275..5a6bb8a9 100644 --- a/fly.toml +++ b/fly.toml @@ -1,6 +1,6 @@ # See https://fly.io/docs/reference/configuration/ for information about how to use this file. -app = 'honcho' +app = 'honcho-patched-v1-test' primary_region = 'ewr' kill_signal = 'SIGINT' kill_timeout = '5s' diff --git a/src/agent.py b/src/agent.py index eafc36b3..8ee03157 100644 --- a/src/agent.py +++ b/src/agent.py @@ -49,7 +49,8 @@ class Dialectic: self.user_representation = user_representation self.chat_history = chat_history self.client = ModelClient( - provider=DEF_DIALECTIC_PROVIDER, model=DEF_DIALECTIC_MODEL + provider=DEF_DIALECTIC_PROVIDER, + model=DEF_DIALECTIC_MODEL, ) self.system_prompt = """You are operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, you'll receive: 1) previously collected psychological context about the user that I've maintained, 2) a series of long-term facts about the user, and 3) their current conversation/interaction from the requesting application. Your goal is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. You are encouraged to provide any context from the provided resources that helps provide a more complete or nuanced understanding of the user, as long as it is somewhat relevant to the query. If the context provided doesn't help address the query, write absolutely NOTHING but "None".""" @@ -76,11 +77,14 @@ class Dialectic: # Create a properly formatted message message: dict[str, Any] = {"role": "user", "content": prompt} - # Generate the response + # Generate the response with fallback logger.debug("Calling model for generation") model_start = asyncio.get_event_loop().time() + response = await self.client.generate( - messages=[message], system=self.system_prompt, max_tokens=1000 + messages=[message], + system=self.system_prompt, + max_tokens=1000, ) model_time = asyncio.get_event_loop().time() - model_start logger.debug( diff --git a/src/crud.py b/src/crud.py index b5ac4f60..44aecd76 100644 --- a/src/crud.py +++ b/src/crud.py @@ -1,3 +1,5 @@ +import os + from collections.abc import Sequence from logging import getLogger from typing import Optional @@ -19,7 +21,7 @@ from .exceptions import ( load_dotenv(override=True) -openai_client = AsyncOpenAI() +openai_client = AsyncOpenAI(base_url=os.getenv("OPENAI_COMPATIBLE_BASE_URL"), api_key=os.getenv("OPENAI_COMPATIBLE_API_KEY")) logger = getLogger(__name__) @@ -1346,7 +1348,7 @@ async def query_documents( top_k: int = 5, ) -> Sequence[models.Document]: # Using async client with await - response = await openai_client.embeddings.create( + response = await openai_client.embeddings.create( model="text-embedding-3-small", input=query ) embedding_query = response.data[0].embedding @@ -1402,7 +1404,7 @@ async def create_document( ) # Using async client with await - response = await openai_client.embeddings.create( + response = await openai_client.embeddings.create( input=document.content, model="text-embedding-3-small" ) @@ -1460,7 +1462,7 @@ async def update_document( if document.content is not None: honcho_document.content = document.content # Using async client with await - response = await openai_client.embeddings.create( + response = await openai_client.embeddings.create( input=document.content, model="text-embedding-3-small" ) embedding = response.data[0].embedding @@ -1531,7 +1533,7 @@ async def get_duplicate_documents( """ # Get embedding for the content # Using async client with await - response = await openai_client.embeddings.create( + response = await openai_client.embeddings.create( input=content, model="text-embedding-3-small" ) embedding = response.data[0].embedding diff --git a/src/db.py b/src/db.py index ee6d544e..f7dad31e 100644 --- a/src/db.py +++ b/src/db.py @@ -5,6 +5,7 @@ from dotenv import load_dotenv from sqlalchemy import MetaData, create_engine, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base +from sqlalchemy.pool import NullPool load_dotenv() @@ -17,12 +18,13 @@ engine = create_async_engine( os.environ["CONNECTION_URI"], connect_args=connect_args, echo=os.getenv("SQL_DEBUG", "false").lower() == "true", # Only enable in debug mode - pool_pre_ping=True, - pool_size=10, - max_overflow=20, - pool_timeout=30, - pool_recycle=300, # Recycle connections after 5 minutes - pool_use_lifo=True, # Use last-in-first-out (LIFO) to prevent connection spread + poolclass=NullPool, + # pool_pre_ping=True, + # pool_size=10, + # max_overflow=20, + # pool_timeout=30, + # pool_recycle=300, # Recycle connections after 5 minutes + # pool_use_lifo=True, # Use last-in-first-out (LIFO) to prevent connection spread ) SessionLocal = async_sessionmaker( diff --git a/src/deriver/tom/conversational.py b/src/deriver/tom/conversational.py index 98b91bfe..4ef774c3 100644 --- a/src/deriver/tom/conversational.py +++ b/src/deriver/tom/conversational.py @@ -1,17 +1,15 @@ import os import sentry_sdk -from anthropic import Anthropic from langfuse.decorators import langfuse_context, observe from sentry_sdk.ai.monitoring import ai_track +from src.utils.model_client import ModelClient, ModelProvider + # Place the code below at the beginning of your application to initialize the tracer -# Initialize the Anthropic client -anthropic = Anthropic( - api_key=os.getenv("ANTHROPIC_API_KEY"), - max_retries=5, -) +# Initialize the unified model client to use OpenAI-compatible endpoint +model_client = ModelClient(provider=ModelProvider.OPENAI) @ai_track("Tom Inference") @@ -71,13 +69,13 @@ async def get_tom_inference_conversational( langfuse_context.update_current_observation( input=messages, model="claude-3-5-sonnet-20240620" ) - message = anthropic.messages.create( - model="claude-3-5-sonnet-20240620", - max_tokens=1000, - temperature=0, + + response = await model_client.generate( messages=messages, + max_tokens=1000, + temperature=0 ) - return message.content[0].text + return response @ai_track("User Representation") @@ -142,10 +140,10 @@ async def get_user_representation_conversational( langfuse_context.update_current_observation( input=messages, model="claude-3-5-sonnet-20240620" ) - message = anthropic.messages.create( - model="claude-3-5-sonnet-20240620", - max_tokens=1000, - temperature=0, + + response = await model_client.generate( messages=messages, + max_tokens=1000, + temperature=0 ) - return message.content[0].text + return response diff --git a/src/utils/model_client.py b/src/utils/model_client.py index f7ed7a88..84002605 100644 --- a/src/utils/model_client.py +++ b/src/utils/model_client.py @@ -13,13 +13,11 @@ from google import genai from google.genai import types as genai_types from langfuse.decorators import langfuse_context, observe -# from openai import AsyncOpenAI from langfuse.openai import AsyncOpenAI # Load environment variables load_dotenv() - # Supported model providers class ModelProvider(str, Enum): ANTHROPIC = "anthropic"