import asyncio import json import os import threading import time import pytest import requests from dotenv import load_dotenv from letta_client import AsyncLetta, MessageCreate, ReasoningMessage, ToolCallMessage from letta_client.core import RequestOptions REASONING_THROTTLE_MS = 100 TEST_USER_MESSAGE = "What products or services does 11x AI sell?" @pytest.fixture(scope="module") def server_url() -> str: """ Provides the URL for the Letta server. If LETTA_SERVER_URL is not set, starts the server in a background thread and polls until it’s accepting connections. """ def _run_server() -> None: load_dotenv() from letta.server.rest_api.app import start_server start_server(debug=True) url: str = os.getenv("LETTA_SERVER_URL", "http://localhost:8283") if not os.getenv("LETTA_SERVER_URL"): thread = threading.Thread(target=_run_server, daemon=True) thread.start() # Poll until the server is up (or timeout) timeout_seconds = 30 deadline = time.time() + timeout_seconds while time.time() < deadline: try: resp = requests.get(url + "/v1/health") if resp.status_code < 500: break except requests.exceptions.RequestException: pass time.sleep(0.1) else: raise RuntimeError(f"Could not reach {url} within {timeout_seconds}s") return url @pytest.fixture(scope="function") def client(server_url: str): """ Creates and returns an asynchronous Letta REST client for testing. """ async_client_instance = AsyncLetta(base_url=server_url) yield async_client_instance async def test_pinecone_tool(client: AsyncLetta) -> None: """ Test the Pinecone tool integration with the Letta client. """ with open("../../scripts/test-afs/knowledge-base.af", "rb") as f: response = await client.agents.import_file(file=f) agent_id = response.agent_ids[0] agent = await client.agents.modify( agent_id=agent_id, tool_exec_environment_variables={ "PINECONE_INDEX_HOST": os.getenv("PINECONE_INDEX_HOST"), "PINECONE_API_KEY": os.getenv("PINECONE_API_KEY"), "PINECONE_NAMESPACE": os.getenv("PINECONE_NAMESPACE"), }, ) last_message = await client.agents.messages.list( agent_id=agent.id, limit=1, ) curr_message_type = None messages = [] reasoning_content = [] last_reasoning_update_ms = 0 tool_call_content = "" tool_return_content = "" summary = None pinecone_results = None queries = [] try: response = client.agents.messages.create_stream( agent_id=agent.id, messages=[ MessageCreate( role="user", content=TEST_USER_MESSAGE, ), ], stream_tokens=True, request_options=RequestOptions( timeout_in_seconds=1000, ), ) async for chunk in response: if chunk.message_type != curr_message_type: messages.append(chunk) curr_message_type = chunk.message_type if curr_message_type == "reasoning_message": reasoning_content = [] if curr_message_type == "tool_call_message": tool_call_content = "" if chunk.message_type == "reasoning_message": now_ms = time.time_ns() // 1_000_000 if now_ms - last_reasoning_update_ms < REASONING_THROTTLE_MS: await asyncio.sleep(REASONING_THROTTLE_MS / 1000) last_reasoning_update_ms = now_ms if len(reasoning_content) == 0: reasoning_content = [chunk.reasoning] else: reasoning_content[-1] += chunk.reasoning message_dict = messages[-1].model_dump() message_dict["reasoning"] = "".join(reasoning_content).strip() messages[-1] = ReasoningMessage(**message_dict) if chunk.message_type == "tool_return_message": tool_return_content += chunk.tool_return if chunk.status == "success": try: if chunk.name == "summarize_pinecone_results": json_response = json.loads(chunk.tool_return) summary = json_response.get("summary", None) pinecone_results = json_response.get("pinecone_results", None) tool_return_content = "" elif chunk.name == "craft_queries": queries.append(chunk.tool_return) tool_return_content = "" except Exception as e: print(f"Error parsing JSON response: {str(e)}. {chunk.tool_return}\n") tool_return_content = "" if chunk.message_type == "tool_call_message": if chunk.tool_call.arguments is not None: tool_call_content += chunk.tool_call.arguments message_dict = messages[-1].model_dump() message_dict["tool_call"]["arguments"] = tool_call_content messages[-1] = ToolCallMessage(**message_dict) except Exception as e: print(f"Failed to fetch knowledge base response: {str(e)}\n") print(tool_call_content) raise e assert len(messages) > 0, "No messages received from the agent." assert len(reasoning_content) > 0, "No reasoning content received from the agent." assert summary is not None, "No summary received from the agent." assert pinecone_results is not None, "No Pinecone results received from the agent." assert len(queries) > 0, "No queries received from the agent." assert messages[-2].message_type == "stop_reason", "Penultimate message in stream must be stop reason." assert messages[-1].message_type == "usage_statistics", "Last message in stream must be usage stats." response_messages_from_stream = [m for m in messages if m.message_type not in ["stop_reason", "usage_statistics"]] response_message_types_from_stream = [m.message_type for m in response_messages_from_stream] messages_from_db = await client.agents.messages.list( agent_id=agent.id, after=last_message[0].id, limit=100, ) response_messages_from_db = [m for m in messages_from_db if m.message_type != "user_message"] response_message_types_from_db = [m.message_type for m in response_messages_from_db] assert len(response_messages_from_stream) == len(response_messages_from_db) assert response_message_types_from_stream == response_message_types_from_db for idx in range(len(response_messages_from_stream)): stream_message = response_messages_from_stream[idx] db_message = response_messages_from_db[idx] assert stream_message.message_type == db_message.message_type assert stream_message.id == db_message.id assert stream_message.otid == db_message.otid if stream_message.message_type == "reasoning_message": assert stream_message.reasoning == db_message.reasoning if stream_message.message_type == "tool_call_message": assert stream_message.tool_call.tool_call_id == db_message.tool_call.tool_call_id assert stream_message.tool_call.name == db_message.tool_call.name if stream_message.tool_call.name == "craft_queries": assert "queries" in stream_message.tool_call.arguments assert "queries" in db_message.tool_call.arguments if stream_message.tool_call.name == "search_and_store_pinecone_records": assert "query_text" in stream_message.tool_call.arguments assert "query_text" in db_message.tool_call.arguments if stream_message.tool_call.name == "summarize_pinecone_results": assert "summary" in stream_message.tool_call.arguments assert "summary" in db_message.tool_call.arguments assert "inner_thoughts" not in stream_message.tool_call.arguments assert "inner_thoughts" not in db_message.tool_call.arguments if stream_message.message_type == "tool_return_message": assert stream_message.tool_return == db_message.tool_return await client.agents.delete(agent_id=agent.id)