solana-agent 1.4.4__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: solana-agent
3
- Version: 1.4.4
3
+ Version: 2.0.0
4
4
  Summary: Build self-learning AI Agents
5
5
  License: MIT
6
6
  Keywords: ai,openai,ai agents
@@ -16,8 +16,7 @@ Classifier: Programming Language :: Python :: 3.12
16
16
  Classifier: Programming Language :: Python :: 3.13
17
17
  Classifier: Programming Language :: Python :: 3 :: Only
18
18
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
- Requires-Dist: cohere (>=5.13.12,<6.0.0)
20
- Requires-Dist: openai (>=1.63.2,<2.0.0)
19
+ Requires-Dist: openai (>=1.64.0,<2.0.0)
21
20
  Requires-Dist: pandas (>=2.2.3,<3.0.0)
22
21
  Requires-Dist: pinecone (>=6.0.1,<7.0.0)
23
22
  Requires-Dist: pydantic (>=2.10.6,<3.0.0)
@@ -58,16 +57,15 @@ Unlike traditional AI assistants that forget conversations after each session, S
58
57
  - Real-time voice-to-voice conversations
59
58
 
60
59
  🧠 **Memory System and Extensibility**
61
- - Advanced AI memory combining conversational context, conversational facts, knowledge base, file search, and parallel tool calling
60
+ - Advanced AI memory combining conversational context, knowledge base, and parallel tool calling
62
61
  - Create custom tools for extending the Agent's capabilities like further API integrations
63
62
 
64
63
  🔍 **Multi-Source Search and Reasoning**
65
64
  - Internet search via Perplexity
66
65
  - X (Twitter) search using Grok
67
- - Conversational fact search powered by Zep
66
+ - Conversational memory powered by Zep
68
67
  - Conversational message history using MongoDB (on-prem or hosted)
69
- - Knowledge Base (KB) using Pinecone with reranking by Cohere - available globally or user-specific
70
- - File uploading and searching using OpenAI like for PDFs
68
+ - Knowledge Base (KB) using Pinecone with reranking - available globally or user-specific
71
69
  - Upload CSVs to be processed into summary reports and stored in the Knowledge Base (KB) using Gemini
72
70
  - Comprehensive reasoning combining multiple data sources
73
71
 
@@ -80,7 +78,7 @@ Unlike traditional AI assistants that forget conversations after each session, S
80
78
  - Persistent cross-session knowledge retention
81
79
  - Automatic self-learning from conversations
82
80
  - Knowledge Base to add domain specific knowledge
83
- - File uploads to perform document context search
81
+ - CSV file uploads to perform document context search
84
82
 
85
83
  🏢 **Enterprise Focus**
86
84
  - Production-ready out of the box in a few lines of code
@@ -29,16 +29,15 @@ Unlike traditional AI assistants that forget conversations after each session, S
29
29
  - Real-time voice-to-voice conversations
30
30
 
31
31
  🧠 **Memory System and Extensibility**
32
- - Advanced AI memory combining conversational context, conversational facts, knowledge base, file search, and parallel tool calling
32
+ - Advanced AI memory combining conversational context, knowledge base, and parallel tool calling
33
33
  - Create custom tools for extending the Agent's capabilities like further API integrations
34
34
 
35
35
  🔍 **Multi-Source Search and Reasoning**
36
36
  - Internet search via Perplexity
37
37
  - X (Twitter) search using Grok
38
- - Conversational fact search powered by Zep
38
+ - Conversational memory powered by Zep
39
39
  - Conversational message history using MongoDB (on-prem or hosted)
40
- - Knowledge Base (KB) using Pinecone with reranking by Cohere - available globally or user-specific
41
- - File uploading and searching using OpenAI like for PDFs
40
+ - Knowledge Base (KB) using Pinecone with reranking - available globally or user-specific
42
41
  - Upload CSVs to be processed into summary reports and stored in the Knowledge Base (KB) using Gemini
43
42
  - Comprehensive reasoning combining multiple data sources
44
43
 
@@ -51,7 +50,7 @@ Unlike traditional AI assistants that forget conversations after each session, S
51
50
  - Persistent cross-session knowledge retention
52
51
  - Automatic self-learning from conversations
53
52
  - Knowledge Base to add domain specific knowledge
54
- - File uploads to perform document context search
53
+ - CSV file uploads to perform document context search
55
54
 
56
55
  🏢 **Enterprise Focus**
57
56
  - Production-ready out of the box in a few lines of code
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "solana-agent"
3
- version = "1.4.4"
3
+ version = "2.0.0"
4
4
  description = "Build self-learning AI Agents"
5
5
  authors = ["Bevan Hunt <bevan@bevanhunt.com>"]
6
6
  license = "MIT"
@@ -18,13 +18,12 @@ python_paths = [".", "tests"]
18
18
 
19
19
  [tool.poetry.dependencies]
20
20
  python = ">=3.9,<4.0"
21
- openai = "^1.63.2"
21
+ openai = "^1.64.0"
22
22
  pydantic = "^2.10.6"
23
23
  pymongo = "^4.11.1"
24
24
  zep-cloud = "^2.4.0"
25
25
  requests = "^2.32.3"
26
26
  pinecone = "^6.0.1"
27
- cohere = "^5.13.12"
28
27
  pandas = "^2.2.3"
29
28
 
30
29
  [build-system]
@@ -1,16 +1,12 @@
1
1
  import asyncio
2
2
  import datetime
3
3
  import json
4
- from typing import AsyncGenerator, Literal, Optional, Dict, Any, Callable
4
+ from typing import AsyncGenerator, List, Literal, Dict, Any, Callable
5
5
  import uuid
6
- import cohere
7
6
  import pandas as pd
8
7
  from pydantic import BaseModel
9
8
  from pymongo import MongoClient
10
9
  from openai import OpenAI
11
- from openai import AssistantEventHandler
12
- from openai.types.beta.threads import TextDelta, Text
13
- from typing_extensions import override
14
10
  import inspect
15
11
  import requests
16
12
  from zep_cloud.client import AsyncZep
@@ -19,88 +15,46 @@ from zep_cloud.types import Message
19
15
  from pinecone import Pinecone
20
16
 
21
17
 
22
- class EventHandler(AssistantEventHandler):
23
- def __init__(self, tool_handlers, ai_instance):
24
- super().__init__()
25
- self._tool_handlers = tool_handlers
26
- self._ai_instance = ai_instance
27
-
28
- @override
29
- def on_text_delta(self, delta: TextDelta, snapshot: Text):
30
- asyncio.create_task(
31
- self._ai_instance._accumulated_value_queue.put(delta.value))
32
-
33
- @override
34
- def on_event(self, event):
35
- if event.event == "thread.run.requires_action":
36
- run_id = event.data.id
37
- self._ai_instance._handle_requires_action(event.data, run_id)
38
-
39
-
40
- class ToolConfig(BaseModel):
41
- name: str
42
- description: str
43
- parameters: Dict[str, Any]
18
+ class DocumentModel(BaseModel):
19
+ id: str
20
+ text: str
44
21
 
45
22
 
46
23
  class MongoDatabase:
47
24
  def __init__(self, db_url: str, db_name: str):
48
25
  self._client = MongoClient(db_url)
49
26
  self.db = self._client[db_name]
50
- self._threads = self.db["threads"]
51
27
  self.messages = self.db["messages"]
52
28
  self.kb = self.db["kb"]
53
- self.vector_stores = self.db["vector_stores"]
54
- self.files = self.db["files"]
55
-
56
- def save_thread_id(self, user_id: str, thread_id: str):
57
- self._threads.insert_one({"thread_id": thread_id, "user_id": user_id})
58
-
59
- def get_thread_id(self, user_id: str) -> Optional[str]:
60
- document = self._threads.find_one({"user_id": user_id})
61
- return document["thread_id"] if document else None
62
29
 
63
30
  def save_message(self, user_id: str, metadata: Dict[str, Any]):
64
31
  metadata["user_id"] = user_id
65
32
  self.messages.insert_one(metadata)
66
33
 
67
- def delete_all_threads(self):
68
- self._threads.delete_many({})
69
-
70
34
  def clear_user_history(self, user_id: str):
71
35
  self.messages.delete_many({"user_id": user_id})
72
- self._threads.delete_one({"user_id": user_id})
73
-
74
- def add_document_to_kb(self, id: str, namespace: str, document: str):
75
- storage = {}
76
- storage["namespace"] = namespace
77
- storage["reference"] = id
78
- storage["document"] = document
79
- storage["timestamp"] = datetime.datetime.now(datetime.timezone.utc)
80
- self.kb.insert_one(storage)
81
-
82
- def get_vector_store_id(self) -> str | None:
83
- document = self.vector_stores.find_one()
84
- return document["vector_store_id"] if document else None
85
36
 
86
- def save_vector_store_id(self, vector_store_id: str):
87
- self.vector_stores.insert_one({"vector_store_id": vector_store_id})
37
+ def add_documents_to_kb(self, namespace: str, documents: List[DocumentModel]):
38
+ for document in documents:
39
+ storage = {}
40
+ storage["namespace"] = namespace
41
+ storage["reference"] = document.id
42
+ storage["document"] = document.text
43
+ storage["timestamp"] = datetime.datetime.now(datetime.timezone.utc)
44
+ self.kb.insert_one(storage)
88
45
 
89
- def delete_vector_store_id(self, vector_store_id: str):
90
- self.vector_stores.delete_one({"vector_store_id": vector_store_id})
91
-
92
- def add_file(self, file_id: str):
93
- self.files.insert_one({"file_id": file_id})
94
-
95
- def delete_file(self, file_id: str):
96
- self.files.delete_one({"file_id": file_id})
46
+ def list_documents_in_kb(self, namespace: str) -> List[DocumentModel]:
47
+ documents = self.kb.find({"namespace": namespace})
48
+ return [
49
+ DocumentModel(id=doc["reference"], text=doc["document"])
50
+ for doc in documents
51
+ ]
97
52
 
98
53
 
99
54
  class AI:
100
55
  def __init__(
101
56
  self,
102
57
  openai_api_key: str,
103
- name: str,
104
58
  instructions: str,
105
59
  database: Any,
106
60
  zep_api_key: str = None,
@@ -108,74 +62,46 @@ class AI:
108
62
  grok_api_key: str = None,
109
63
  pinecone_api_key: str = None,
110
64
  pinecone_index_name: str = None,
111
- cohere_api_key: str = None,
112
- cohere_model: Literal["rerank-v3.5"] = "rerank-v3.5",
65
+ pinecone_embed_model: Literal["llama-text-embed-v2"] = "llama-text-embed-v2",
113
66
  gemini_api_key: str = None,
114
- code_interpreter: bool = False,
115
- file_search: bool = False,
116
- openai_assistant_model: Literal["gpt-4o-mini",
117
- "gpt-4o"] = "gpt-4o-mini",
118
- openai_embedding_model: Literal[
119
- "text-embedding-3-small", "text-embedding-3-large"
120
- ] = "text-embedding-3-large",
67
+ openai_base_url: str = None,
68
+ openai_model: str = "gpt-4o-mini",
121
69
  ):
122
- """Initialize a new AI assistant with memory and tool integration capabilities.
70
+ """Initialize a new AI assistant instance.
123
71
 
124
72
  Args:
125
73
  openai_api_key (str): OpenAI API key for core AI functionality
126
- name (str): Name identifier for the assistant
127
74
  instructions (str): Base behavioral instructions for the AI
128
75
  database (Any): Database instance for message/thread storage
129
- zep_api_key (str, optional): API key for Zep memory integration. Defaults to None
76
+ zep_api_key (str, optional): API key for Zep memory storage. Defaults to None
130
77
  perplexity_api_key (str, optional): API key for Perplexity search. Defaults to None
131
78
  grok_api_key (str, optional): API key for X/Twitter search via Grok. Defaults to None
132
79
  pinecone_api_key (str, optional): API key for Pinecone. Defaults to None
133
80
  pinecone_index_name (str, optional): Name of the Pinecone index. Defaults to None
134
- cohere_api_key (str, optional): API key for Cohere search. Defaults to None
135
- cohere_model (Literal["rerank-v3.5"], optional): Cohere model for reranking. Defaults to "rerank-v3.5"
81
+ pinecone_embed_model (Literal["llama-text-embed-v2"], optional): Pinecone embedding model. Defaults to "llama-text-embed-v2"
136
82
  gemini_api_key (str, optional): API key for Gemini search. Defaults to None
137
- code_interpreter (bool, optional): Enable code interpretation. Defaults to False
138
- file_search (bool, optional): Enable file search tool. Defaults to False
139
- openai_assistant_model (Literal["gpt-4o-mini", "gpt-4o"], optional): OpenAI model for assistant. Defaults to "gpt-4o-mini"
140
- openai_embedding_model (Literal["text-embedding-3-small", "text-embedding-3-large"], optional): OpenAI model for text embedding. Defaults to "text-embedding-3-large"
83
+ openai_base_url (str, optional): Base URL for OpenAI API. Defaults to None
84
+ openai_model (str, optional): OpenAI model to use. Defaults to "gpt-4o-mini"
141
85
 
142
86
  Example:
143
87
  ```python
144
88
  ai = AI(
145
89
  openai_api_key="your-key",
146
- name="Assistant",
147
90
  instructions="Be helpful and concise",
148
- database=MongoDatabase("mongodb://localhost", "ai_db"),
91
+ database=MongoDatabase("mongodb://localhost", "ai_db")
149
92
  )
150
93
  ```
151
94
  Notes:
152
95
  - Requires valid OpenAI API key for core functionality
96
+ - Supports any OpenAI compatible model for conversation
97
+ - Requires valid Zep API key for memory storage
153
98
  - Database instance for storing messages and threads
154
- - Optional integrations for Zep, Perplexity, Pinecone, Cohere, Gemini, and Grok
155
- - Supports code interpretation and custom tool functions
99
+ - Optional integrations for Perplexity, Pinecone, Gemini, and Grok
156
100
  - You must create the Pinecone index in the dashboard before using it
157
101
  """
158
- self._client = OpenAI(api_key=openai_api_key)
159
- self._name = name
102
+ self._client = OpenAI(api_key=openai_api_key, base_url=openai_base_url) if openai_base_url else OpenAI(
103
+ api_key=openai_api_key)
160
104
  self._instructions = instructions
161
- self._openai_assistant_model = openai_assistant_model
162
- self._openai_embedding_model = openai_embedding_model
163
- self._file_search = file_search
164
- if file_search:
165
- self._tools = (
166
- [
167
- {"type": "code_interpreter"},
168
- {"type": "file_search"},
169
- ]
170
- if code_interpreter
171
- else [{"type": "file_search"}]
172
- )
173
- else:
174
- self._tools = [{"type": "code_interpreter"}
175
- ] if code_interpreter else []
176
-
177
- self._tool_handlers = {}
178
- self._assistant_id = None
179
105
  self._database: MongoDatabase = database
180
106
  self._accumulated_value_queue = asyncio.Queue()
181
107
  self._zep = AsyncZep(api_key=zep_api_key) if zep_api_key else None
@@ -187,94 +113,22 @@ class AI:
187
113
  Pinecone(api_key=pinecone_api_key) if pinecone_api_key else None
188
114
  )
189
115
  self._pinecone_index_name = pinecone_index_name if pinecone_index_name else None
116
+ self._pinecone_embedding_model = pinecone_embed_model
190
117
  self.kb = (
191
118
  self._pinecone.Index(
192
119
  self._pinecone_index_name) if self._pinecone else None
193
120
  )
194
- self._co = cohere.ClientV2(
195
- api_key=cohere_api_key) if cohere_api_key else None
196
- self._co_model = cohere_model if cohere_api_key else None
121
+ self._openai_base_url = openai_base_url
122
+ self._openai_model = openai_model
123
+ self._tools = []
197
124
 
198
125
  async def __aenter__(self):
199
- assistants = self._client.beta.assistants.list()
200
- existing_assistant = next(
201
- (a for a in assistants if a.name == self._name), None)
202
-
203
- if existing_assistant:
204
- self._assistant_id = existing_assistant.id
205
- else:
206
- self._assistant_id = self._client.beta.assistants.create(
207
- name=self._name,
208
- instructions=self._instructions,
209
- tools=self._tools,
210
- model=self._openai_assistant_model,
211
- ).id
212
- self._database.delete_all_threads()
213
- if self._file_search:
214
- vectore_store_id = self._database.get_vector_store_id()
215
- if vectore_store_id:
216
- self._vector_store = self._client.beta.vector_stores.retrieve(
217
- vector_store_id=vectore_store_id
218
- )
219
- else:
220
- uid = uuid.uuid4().hex
221
- self._vector_store = self._client.beta.vector_stores.create(
222
- name=uid)
223
- self._database.save_vector_store_id(self._vector_store.id)
224
- self._client.beta.assistants.update(
225
- assistant_id=self._assistant_id,
226
- tool_resources={
227
- "file_search": {"vector_store_ids": [self._vector_store.id]}
228
- },
229
- )
230
126
  return self
231
127
 
232
128
  async def __aexit__(self, exc_type, exc_val, exc_tb):
233
129
  # Perform any cleanup actions here
234
130
  pass
235
131
 
236
- async def _create_thread(self, user_id: str) -> str:
237
- thread_id = self._database.get_thread_id(user_id)
238
-
239
- if thread_id is None:
240
- thread = self._client.beta.threads.create()
241
- thread_id = thread.id
242
- self._database.save_thread_id(user_id, thread_id)
243
- if self._zep:
244
- try:
245
- await self._zep.user.add(user_id=user_id)
246
- except Exception:
247
- pass
248
- try:
249
- await self._zep.memory.add_session(
250
- user_id=user_id, session_id=user_id
251
- )
252
- except Exception:
253
- pass
254
-
255
- return thread_id
256
-
257
- async def _cancel_run(self, thread_id: str, run_id: str):
258
- try:
259
- self._client.beta.threads.runs.cancel(
260
- thread_id=thread_id, run_id=run_id)
261
- except Exception as e:
262
- print(f"Error cancelling run: {e}")
263
-
264
- async def _get_active_run(self, thread_id: str) -> Optional[str]:
265
- runs = self._client.beta.threads.runs.list(
266
- thread_id=thread_id, limit=1)
267
- for run in runs:
268
- if run.status in ["in_progress"]:
269
- return run.id
270
- return None
271
-
272
- async def _get_run_status(self, thread_id: str, run_id: str) -> str:
273
- run = self._client.beta.threads.runs.retrieve(
274
- thread_id=thread_id, run_id=run_id
275
- )
276
- return run.status
277
-
278
132
  def csv_to_text(self, file, filename: str) -> str:
279
133
  """Convert a CSV file to a Markdown table text format optimized for LLM ingestion.
280
134
 
@@ -368,20 +222,20 @@ class AI:
368
222
  self,
369
223
  file,
370
224
  filename: str,
371
- prompt: str = "Summarize the markdown table into a report, include important metrics and totals.",
225
+ id: str = uuid.uuid4().hex,
226
+ prompt: str = "Summarize the table into a report, include important metrics and totals.",
372
227
  namespace: str = "global",
373
- model: Literal["gemini-2.0-flash",
374
- "gemini-1.5-pro"] = "gemini-1.5-pro",
228
+ model: Literal["gemini-2.0-flash"] = "gemini-2.0-flash",
375
229
  ):
376
230
  """Upload and process a CSV file into the knowledge base with AI summarization.
377
231
 
378
232
  Args:
379
233
  file (BinaryIO): The CSV file to upload and process
380
234
  filename (str): The name of the CSV file
381
- prompt (str, optional): Custom prompt for summarization. Defaults to "Summarize the markdown table into a report, include important metrics and totals."
235
+ id (str, optional): Unique identifier for the document. Defaults to a random UUID.
236
+ prompt (str, optional): Custom prompt for summarization. Defaults to "Summarize the table into a report, include important metrics and totals."
382
237
  namespace (str, optional): Knowledge base namespace. Defaults to "global".
383
- model (Literal["gemini-2.0-flash", "gemini-1.5-pro"], optional):
384
- Gemini model for summarization. Defaults to "gemini-1.5-pro"
238
+ model (Literal["gemini-2.0-flash"], optional): Gemini model for summarization. Defaults to "gemini-2.0-flash".
385
239
 
386
240
  Example:
387
241
  ```python
@@ -393,123 +247,32 @@ class AI:
393
247
 
394
248
  Note:
395
249
  - Converts CSV to Markdown table format
396
- - Uses Gemini AI to generate a summary
250
+ - Uses Gemini AI to generate a summary - total of 1M context tokens
397
251
  - Stores summary in Pinecone knowledge base
398
252
  - Requires configured Pinecone index
399
253
  - Supports custom prompts for targeted summaries
400
254
  """
401
255
  csv_text = self.csv_to_text(file, filename)
402
- print(csv_text)
403
256
  document = self.summarize(csv_text, prompt, model)
404
- print(document)
405
- self.add_document_to_kb(document=document, namespace=namespace)
406
-
407
- def delete_vector_store_and_files(self):
408
- """Delete the OpenAI vector store and files.
409
-
410
- Example:
411
- ```python
412
- ai.delete_vector_store()
413
- ```
414
-
415
- Note:
416
- - Requires file_search=True in AI initialization
417
- - Deletes the vector store and all associated files
418
- """
419
- vector_store_id = self._database.get_vector_store_id()
420
- if vector_store_id:
421
- self._client.beta.vector_stores.delete(vector_store_id)
422
- self._database.delete_vector_store_id(vector_store_id)
423
- for file in self._database.files.find().to_list():
424
- self._client.files.delete(file["file_id"])
425
- self._database.delete_file(file["file_id"])
426
-
427
- def max_files(self) -> bool:
428
- """Check if the OpenAI vector store has reached its maximum file capacity.
429
-
430
- Returns:
431
- bool: True if file count is at maximum (10,000), False otherwise
432
-
433
- Example:
434
- ```python
435
- if ai.max_files():
436
- print("Vector store is full")
437
- else:
438
- print("Can still add more files")
439
- ```
440
-
441
- Note:
442
- - Requires file_search=True in AI initialization
443
- - OpenAI vector stores have a 10,000 file limit
444
- - Returns False if vector store is not configured
445
- """
446
- self._vector_store.file_counts.completed == 10000
447
-
448
- def file_count(self) -> int:
449
- """Get the total number of files processed in the OpenAI vector store.
450
-
451
- Returns:
452
- int: Number of successfully processed files in the vector store
453
-
454
- Example:
455
- ```python
456
- count = ai.file_count()
457
- print(f"Processed {count} files")
458
- # Returns: "Processed 5 files"
459
- ```
460
-
461
- Note:
462
- - Requires file_search=True in AI initialization
463
- - Only counts successfully processed files
464
- - Returns 0 if vector store is not configured
465
- """
466
- self._vector_store.file_counts.completed
467
-
468
- def add_file(
469
- self,
470
- filename: str,
471
- file_stream: bytes,
472
- ) -> Literal["in_progress", "completed", "cancelled", "failed"]:
473
- """Upload and process a file in the OpenAI vector store.
474
-
475
- Args:
476
- filename (str): Name of the file to upload
477
- file_stream (bytes): Raw bytes of the file to upload
478
-
479
- Returns:
480
- Literal["in_progress", "completed", "cancelled", "failed"]: Status of file processing
481
-
482
- Example:
483
- ```python
484
- with open('document.pdf', 'rb') as f:
485
- status = ai.add_file(f.filename, f.read())
486
- if status == "completed":
487
- print("File processed successfully")
488
- ```
489
-
490
- Note:
491
- - Requires file_search=True in AI initialization
492
- - Files are vectorized for semantic search
493
- - Maximum file size: 512MB
494
- - Maximum 10,000 files per vector store
495
- - Processing may take a few seconds to minutes
496
- """
497
- vector_store_id = self._database.get_vector_store_id()
498
- file = self._client.files.create(
499
- file=(filename, file_stream), purpose="assistants"
257
+ self.add_documents_to_kb(
258
+ documents=[DocumentModel(id=id, text=document)], namespace=namespace
500
259
  )
501
- file_batch = self._client.beta.vector_stores.files.create_and_poll(
502
- vector_store_id=vector_store_id, file_id=file.id
503
- )
504
- self._database.add_file(file.id)
505
- return file_batch.status
506
260
 
507
- def search_kb(self, query: str, namespace: str = "global", limit: int = 3) -> str:
508
- """Search Pinecone knowledge base using OpenAI embeddings.
261
+ def search_kb(
262
+ self,
263
+ query: str,
264
+ namespace: str = "global",
265
+ rerank_model: Literal["cohere-rerank-3.5"] = "cohere-rerank-3.5",
266
+ inner_limit: int = 10,
267
+ limit: int = 3,
268
+ ) -> str:
269
+ """Search Pinecone knowledge base.
509
270
 
510
271
  Args:
511
272
  query (str): Search query to find relevant documents
512
273
  namespace (str, optional): Namespace of the Pinecone to search. Defaults to "global".
274
+ rerank_model (Literal["cohere-rerank-3.5"], optional): Rerank model on Pinecone. Defaults to "cohere-rerank-3.5".
275
+ inner_limit (int, optional): Maximum number of results to rerank. Defaults to 10.
513
276
  limit (int, optional): Maximum number of results to return. Defaults to 3.
514
277
 
515
278
  Returns:
@@ -517,25 +280,23 @@ class AI:
517
280
 
518
281
  Example:
519
282
  ```python
520
- results = ai.search_kb("user123", "machine learning basics")
283
+ results = ai.search_kb("machine learning basics", "user123")
521
284
  # Returns: '["Document 1", "Document 2", ...]'
522
285
  ```
523
286
 
524
287
  Note:
525
288
  - Requires configured Pinecone index
526
- - Uses OpenAI embeddings for semantic search
527
- - Returns JSON-serialized Pinecone match metadata results
528
289
  - Returns error message string if search fails
529
- - Optionally reranks results using Cohere API
530
290
  """
531
291
  try:
532
- response = self._client.embeddings.create(
533
- input=query,
534
- model=self._openai_embedding_model,
292
+ embedding = self._pinecone.inference.embed(
293
+ model=self._pinecone_embedding_model,
294
+ inputs=[query],
295
+ parameters={"input_type": "query"},
535
296
  )
536
297
  search_results = self.kb.query(
537
- vector=response.data[0].embedding,
538
- top_k=10,
298
+ vector=embedding[0].values,
299
+ top_k=inner_limit,
539
300
  include_metadata=False,
540
301
  include_values=False,
541
302
  namespace=namespace,
@@ -548,62 +309,88 @@ class AI:
548
309
  for id in ids:
549
310
  document = self._database.kb.find_one({"reference": id})
550
311
  docs.append(document["document"])
551
- if self._co:
552
- try:
553
- response = self._co.rerank(
554
- model=self._co_model,
555
- query=query,
556
- documents=docs,
557
- top_n=limit,
558
- )
559
- reranked_docs = response.results
560
- new_docs = []
561
- for doc in reranked_docs:
562
- new_docs.append(docs[doc.index])
563
- return json.dumps(new_docs)
564
- except Exception:
565
- return json.dumps(docs[:limit])
566
- else:
312
+ try:
313
+ reranked_docs = self._pinecone.inference.rerank(
314
+ model=rerank_model,
315
+ query=query,
316
+ documents=docs,
317
+ top_n=limit,
318
+ )
319
+ new_docs = []
320
+ for doc in reranked_docs.data:
321
+ new_docs.append(docs[doc.index])
322
+ return json.dumps(new_docs)
323
+ except Exception:
567
324
  return json.dumps(docs[:limit])
568
325
  except Exception as e:
569
326
  return f"Failed to search KB. Error: {e}"
570
327
 
571
- def add_document_to_kb(
328
+ def list_documents_in_kb(self, namespace: str = "global") -> List[DocumentModel]:
329
+ """List all documents stored in the Pinecone knowledge base.
330
+
331
+ Args:
332
+ namespace (str, optional): Namespace of the Pinecone index to search. Defaults to "global".
333
+
334
+ Returns:
335
+ List[DocumentModel]: List of documents stored in the knowledge base
336
+
337
+ Example:
338
+ ```python
339
+ documents = ai.list_documents_in_kb("user123")
340
+ for doc in documents:
341
+ print(doc)
342
+ # Returns: "Document 1", "Document 2", ...
343
+ ```
344
+
345
+ Note:
346
+ - Requires Pinecone index to be configured
347
+ """
348
+ return self._database.list_documents_in_kb(namespace)
349
+
350
+ def add_documents_to_kb(
572
351
  self,
573
- document: str,
574
- id: str = uuid.uuid4().hex,
352
+ documents: List[DocumentModel],
575
353
  namespace: str = "global",
576
354
  ):
577
- """Add a document to the Pinecone knowledge base with OpenAI embeddings.
355
+ """Add documents to the Pinecone knowledge base.
578
356
 
579
357
  Args:
580
- document (str): Document to add to the knowledge base
581
- id (str, optional): Unique identifier for the document. Defaults to random UUID.
358
+ documents (List[DocumentModel]): List of documents to add to the knowledge base
582
359
  namespace (str): Namespace of the Pinecone index to search. Defaults to "global".
583
360
 
584
361
  Example:
585
362
  ```python
586
- ai.add_document_to_kb("user123 has 4 cats")
363
+ docs = [
364
+ {"id": "doc1", "text": "Document 1"},
365
+ {"id": "doc2", "text": "Document 2"},
366
+ ]
367
+ ai.add_documents_to_kb(docs, "user123")
587
368
  ```
588
369
 
589
370
  Note:
590
371
  - Requires Pinecone index to be configured
591
- - Uses OpenAI embeddings API
592
372
  """
593
- response = self._client.embeddings.create(
594
- input=document,
595
- model=self._openai_embedding_model,
373
+ embeddings = self._pinecone.inference.embed(
374
+ model=self._pinecone_embedding_model,
375
+ inputs=[d.text for d in documents],
376
+ parameters={"input_type": "passage", "truncate": "END"},
596
377
  )
597
- self.kb.upsert(
598
- vectors=[
378
+
379
+ vectors = []
380
+ for d, e in zip(documents, embeddings):
381
+ vectors.append(
599
382
  {
600
- "id": id,
601
- "values": response.data[0].embedding,
383
+ "id": d.id,
384
+ "values": e["values"],
602
385
  }
603
- ],
386
+ )
387
+
388
+ self.kb.upsert(
389
+ vectors=vectors,
604
390
  namespace=namespace,
605
391
  )
606
- self._database.add_document_to_kb(id, namespace, document)
392
+
393
+ self._database.add_documents_to_kb(namespace, documents)
607
394
 
608
395
  def delete_document_from_kb(self, id: str, user_id: str = "global"):
609
396
  """Delete a document from the Pinecone knowledge base.
@@ -644,54 +431,37 @@ class AI:
644
431
  "%Y-%m-%d %H:%M:%S %Z"
645
432
  )
646
433
 
647
- # search facts tool - has to be sync
648
- def search_facts(
434
+ # has to be sync for tool
435
+ def get_memory_context(
649
436
  self,
650
437
  user_id: str,
651
- query: str,
652
- limit: int = 10,
653
438
  ) -> str:
654
- """Search stored conversation facts using Zep memory integration.
439
+ """Retrieve contextual memory for a specific user from Zep memory storage.
655
440
 
656
441
  Args:
657
- user_id (str): Unique identifier for the user
658
- query (str): Search query to find relevant facts
659
- limit (int, optional): Maximum number of facts to return. Defaults to 10.
442
+ user_id (str): Unique identifier for the user whose memory context to retrieve
660
443
 
661
444
  Returns:
662
- str: JSON string of matched facts or error message
445
+ str: User's memory context or error message if retrieval fails
663
446
 
664
447
  Example:
665
448
  ```python
666
- facts = ai.search_facts(
667
- user_id="user123",
668
- query="How many cats do I have?"
669
- )
670
- # Returns: [{"fact": "user123 has 4 cats", "timestamp": "2022-01-01T12:00:00Z"}]
449
+ context = ai.get_memory_context("user123")
450
+ print(context)
451
+ # Returns: "User previously mentioned having 3 dogs and living in London"
671
452
  ```
672
453
 
673
454
  Note:
674
- Requires Zep integration to be configured with valid API key.
675
- This is a synchronous tool method required for OpenAI function calling.
455
+ - This is a synchronous tool method required for OpenAI function calling
456
+ - Requires Zep integration to be configured with valid API key
457
+ - Returns error message if Zep is not configured or retrieval fails
458
+ - Useful for maintaining conversation context across sessions
676
459
  """
677
- if self._sync_zep:
678
- try:
679
- facts = []
680
- results = self._sync_zep.memory.search_sessions(
681
- user_id=user_id,
682
- session_ids=[user_id],
683
- text=query,
684
- limit=limit,
685
- )
686
- for result in results.results:
687
- fact = result.fact.fact
688
- timestamp = result.fact.created_at
689
- facts.append({"fact": fact, "timestamp": timestamp})
690
- return json.dumps(facts)
691
- except Exception as e:
692
- return f"Failed to search facts. Error: {e}"
693
- else:
694
- return "Zep integration not configured."
460
+ try:
461
+ memory = self._sync_zep.memory.get(session_id=user_id)
462
+ return memory.context
463
+ except Exception:
464
+ return ""
695
465
 
696
466
  # search internet tool - has to be sync
697
467
  def search_internet(
@@ -769,7 +539,6 @@ class AI:
769
539
  prompt: str = "You combine the data with your reasoning to answer the query.",
770
540
  use_perplexity: bool = True,
771
541
  use_grok: bool = True,
772
- use_facts: bool = True,
773
542
  use_kb: bool = True,
774
543
  perplexity_model: Literal[
775
544
  "sonar", "sonar-pro", "sonar-reasoning-pro", "sonar-reasoning"
@@ -786,7 +555,6 @@ class AI:
786
555
  prompt (str, optional): Prompt for reasoning. Defaults to "You combine the data with your reasoning to answer the query."
787
556
  use_perplexity (bool, optional): Include Perplexity search results. Defaults to True
788
557
  use_grok (bool, optional): Include X/Twitter search results. Defaults to True
789
- use_facts (bool, optional): Include stored conversation facts. Defaults to True
790
558
  use_kb (bool, optional): Include Pinecone knowledge base search results. Defaults to True
791
559
  perplexity_model (Literal, optional): Perplexity model to use. Defaults to "sonar"
792
560
  openai_model (Literal, optional): OpenAI model for reasoning. Defaults to "o3-mini"
@@ -818,13 +586,7 @@ class AI:
818
586
  kb_results = ""
819
587
  else:
820
588
  kb_results = ""
821
- if use_facts:
822
- try:
823
- facts = self.search_facts(user_id, query)
824
- except Exception:
825
- facts = ""
826
- else:
827
- facts = ""
589
+
828
590
  if use_perplexity:
829
591
  try:
830
592
  search_results = self.search_internet(
@@ -833,6 +595,7 @@ class AI:
833
595
  search_results = ""
834
596
  else:
835
597
  search_results = ""
598
+
836
599
  if use_grok:
837
600
  try:
838
601
  x_search_results = self.search_x(query, grok_model)
@@ -841,6 +604,11 @@ class AI:
841
604
  else:
842
605
  x_search_results = ""
843
606
 
607
+ if self._zep:
608
+ memory = self._sync_zep.memory.get(session_id=user_id)
609
+ else:
610
+ memory = ""
611
+
844
612
  response = self._client.chat.completions.create(
845
613
  model=openai_model,
846
614
  messages=[
@@ -850,7 +618,7 @@ class AI:
850
618
  },
851
619
  {
852
620
  "role": "user",
853
- "content": f"Query: {query}, Facts: {facts}, KB Results: {kb_results}, Internet Search Results: {search_results}, X Search Results: {x_search_results}",
621
+ "content": f"Query: {query}, Memory: {memory}, KB Results: {kb_results}, Internet Search Results: {search_results}, X Search Results: {x_search_results}",
854
622
  },
855
623
  ],
856
624
  )
@@ -916,16 +684,12 @@ class AI:
916
684
  Note:
917
685
  This is an async method and must be awaited.
918
686
  """
919
- try:
920
- await self.delete_assistant_thread(user_id)
921
- except Exception:
922
- pass
923
687
  try:
924
688
  self._database.clear_user_history(user_id)
925
689
  except Exception:
926
690
  pass
927
691
  try:
928
- await self.delete_facts(user_id)
692
+ await self.delete_memory(user_id)
929
693
  except Exception:
930
694
  pass
931
695
 
@@ -941,8 +705,8 @@ class AI:
941
705
  thread_id = self._database.get_thread_id(user_id)
942
706
  await self._client.beta.threads.delete(thread_id=thread_id)
943
707
 
944
- async def delete_facts(self, user_id: str):
945
- """Delete stored conversation facts for a specific user from Zep memory.
708
+ async def delete_memory(self, user_id: str):
709
+ """Delete memory for a specific user from Zep memory.
946
710
 
947
711
  Args:
948
712
  user_id (str): Unique identifier for the user whose facts should be deleted
@@ -973,11 +737,11 @@ class AI:
973
737
  """Process text input and stream AI responses asynchronously.
974
738
 
975
739
  Args:
976
- user_id (str): Unique identifier for the user/conversation
977
- user_text (str): Text input from user to process
740
+ user_id (str): Unique identifier for the user/conversation.
741
+ user_text (str): Text input from user to process.
978
742
 
979
743
  Returns:
980
- AsyncGenerator[str, None]: Stream of response text chunks
744
+ AsyncGenerator[str, None]: Stream of response text chunks (including tool call results).
981
745
 
982
746
  Example:
983
747
  ```python
@@ -986,69 +750,108 @@ class AI:
986
750
  ```
987
751
 
988
752
  Note:
989
- - Maintains conversation thread using OpenAI's thread system
990
- - Stores messages in configured database (MongoDB/SQLite)
991
- - Integrates with Zep memory if configured
992
- - Handles concurrent runs by canceling active ones
993
- - Streams responses for real-time interaction
753
+ - Maintains conversation thread using OpenAI's thread system.
754
+ - Stores messages in configured database (MongoDB/SQLite).
755
+ - Integrates with Zep memory if configured.
756
+ - Supports tool calls by aggregating and executing them as their arguments stream in.
994
757
  """
995
758
  self._accumulated_value_queue = asyncio.Queue()
996
-
997
- thread_id = self._database.get_thread_id(user_id)
998
-
999
- if thread_id is None:
1000
- thread_id = await self._create_thread(user_id)
1001
-
1002
- self._current_thread_id = thread_id
1003
-
1004
- # Check for active runs and cancel if necessary
1005
- active_run_id = await self._get_active_run(thread_id)
1006
- if active_run_id:
1007
- await self._cancel_run(thread_id, active_run_id)
1008
- while await self._get_run_status(thread_id, active_run_id) != "cancelled":
1009
- await asyncio.sleep(0.1)
1010
-
1011
- # Create a message in the thread
1012
- self._client.beta.threads.messages.create(
1013
- thread_id=thread_id,
1014
- role="user",
1015
- content=user_text,
1016
- )
1017
- event_handler = EventHandler(self._tool_handlers, self)
759
+ final_tool_calls = {} # Accumulate tool call deltas
760
+ final_response = ""
1018
761
 
1019
762
  async def stream_processor():
1020
- with self._client.beta.threads.runs.stream(
1021
- thread_id=thread_id,
1022
- assistant_id=self._assistant_id,
1023
- event_handler=event_handler,
1024
- ) as stream:
1025
- stream.until_done()
1026
-
1027
- # Start the stream processor in a separate task
763
+ memory = self.get_memory_context(user_id)
764
+ response = self._client.chat.completions.create(
765
+ model=self._openai_model,
766
+ messages=[
767
+ {
768
+ "role": "system",
769
+ "content": self._instructions,
770
+ },
771
+ {
772
+ "role": "user",
773
+ "content": f"Query: {user_text}, Memory: {memory}",
774
+ },
775
+ ],
776
+ tools=self._tools,
777
+ stream=True,
778
+ )
779
+ for chunk in response:
780
+ delta = chunk.choices[0].delta
781
+
782
+ # Process tool call deltas (if any)
783
+ if delta.tool_calls:
784
+ for tool_call in delta.tool_calls:
785
+ index = tool_call.index
786
+ if tool_call.function.name:
787
+ # Initialize a new tool call record
788
+ final_tool_calls[index] = {
789
+ "name": tool_call.function.name,
790
+ "arguments": tool_call.function.arguments or ""
791
+ }
792
+ elif tool_call.function.arguments:
793
+ # Append additional arguments if provided in subsequent chunks
794
+ final_tool_calls[index]["arguments"] += tool_call.function.arguments
795
+
796
+ try:
797
+ args = json.loads(
798
+ final_tool_calls[index]["arguments"])
799
+ func = getattr(
800
+ self, final_tool_calls[index]["name"])
801
+ # Execute the tool call (synchronously; adjust if async is needed)
802
+ result = func(**args)
803
+ response = self._client.chat.completions.create(
804
+ model=self._openai_model,
805
+ messages=[
806
+ {
807
+ "role": "system",
808
+ "content": self._instructions,
809
+ },
810
+ {
811
+ "role": "user",
812
+ "content": result,
813
+ },
814
+ ],
815
+ tools=self._tools,
816
+ stream=True,
817
+ )
818
+ for chunk in response:
819
+ delta = chunk.choices[0].delta
820
+
821
+ if delta.content is not None:
822
+ await self._accumulated_value_queue.put(delta.content)
823
+ # Remove the cached tool call info so it doesn't block future calls
824
+ del final_tool_calls[index]
825
+ except json.JSONDecodeError:
826
+ # If the accumulated arguments aren't valid yet, wait for more chunks.
827
+ continue
828
+
829
+ # Process regular response content
830
+ if delta.content is not None:
831
+ await self._accumulated_value_queue.put(delta.content)
832
+
833
+ # Start the stream processor as a background task
1028
834
  asyncio.create_task(stream_processor())
1029
835
 
1030
- # Yield values from the queue as they become available
1031
- full_response = ""
836
+ # Yield values from the queue as they become available.
1032
837
  while True:
1033
838
  try:
1034
- value = await asyncio.wait_for(
1035
- self._accumulated_value_queue.get(), timeout=0.1
1036
- )
839
+ value = await asyncio.wait_for(self._accumulated_value_queue.get(), timeout=0.1)
1037
840
  if value is not None:
1038
- full_response += value
841
+ final_response += value
1039
842
  yield value
1040
843
  except asyncio.TimeoutError:
844
+ # Break only if the queue is empty (assuming stream ended)
1041
845
  if self._accumulated_value_queue.empty():
1042
846
  break
1043
847
 
1044
- # Save the message to the database
848
+ # Save the conversation to the database and Zep memory (if configured)
1045
849
  metadata = {
1046
850
  "user_id": user_id,
1047
851
  "message": user_text,
1048
- "response": full_response,
852
+ "response": final_response,
1049
853
  "timestamp": datetime.datetime.now(datetime.timezone.utc),
1050
854
  }
1051
-
1052
855
  self._database.save_message(user_id, metadata)
1053
856
  if self._zep:
1054
857
  messages = [
@@ -1060,7 +863,7 @@ class AI:
1060
863
  Message(
1061
864
  role="assistant",
1062
865
  role_type="assistant",
1063
- content=full_response,
866
+ content=final_response,
1064
867
  ),
1065
868
  ]
1066
869
  await self._zep.memory.add(session_id=user_id, messages=messages)
@@ -1112,56 +915,105 @@ class AI:
1112
915
  - Streams audio response using OpenAI TTS
1113
916
  """
1114
917
 
1115
- # Reset the queue for each new conversation
1116
- self._accumulated_value_queue = asyncio.Queue()
1117
-
1118
- thread_id = self._database.get_thread_id(user_id)
1119
-
1120
- if thread_id is None:
1121
- thread_id = await self._create_thread(user_id)
1122
-
1123
- self._current_thread_id = thread_id
1124
918
  transcript = await self._listen(audio_bytes, input_format)
1125
- event_handler = EventHandler(self._tool_handlers, self)
1126
- self._client.beta.threads.messages.create(
1127
- thread_id=thread_id,
1128
- role="user",
1129
- content=transcript,
1130
- )
919
+ self._accumulated_value_queue = asyncio.Queue()
920
+ final_tool_calls = {} # Accumulate tool call deltas
921
+ final_response = ""
1131
922
 
1132
923
  async def stream_processor():
1133
- with self._client.beta.threads.runs.stream(
1134
- thread_id=thread_id,
1135
- assistant_id=self._assistant_id,
1136
- event_handler=event_handler,
1137
- ) as stream:
1138
- stream.until_done()
1139
-
1140
- # Start the stream processor in a separate task
924
+ memory = self.get_memory_context(user_id)
925
+ response = self._client.chat.completions.create(
926
+ model=self._openai_model,
927
+ messages=[
928
+ {
929
+ "role": "system",
930
+ "content": self._instructions,
931
+ },
932
+ {
933
+ "role": "user",
934
+ "content": f"Query: {transcript}, Memory: {memory}",
935
+ },
936
+ ],
937
+ tools=self._tools,
938
+ stream=True,
939
+ )
940
+ for chunk in response:
941
+ delta = chunk.choices[0].delta
942
+
943
+ # Process tool call deltas (if any)
944
+ if delta.tool_calls:
945
+ for tool_call in delta.tool_calls:
946
+ index = tool_call.index
947
+ if tool_call.function.name:
948
+ # Initialize a new tool call record
949
+ final_tool_calls[index] = {
950
+ "name": tool_call.function.name,
951
+ "arguments": tool_call.function.arguments or ""
952
+ }
953
+ elif tool_call.function.arguments:
954
+ # Append additional arguments if provided in subsequent chunks
955
+ final_tool_calls[index]["arguments"] += tool_call.function.arguments
956
+
957
+ try:
958
+ args = json.loads(
959
+ final_tool_calls[index]["arguments"])
960
+ func = getattr(
961
+ self, final_tool_calls[index]["name"])
962
+ # Execute the tool call (synchronously; adjust if async is needed)
963
+ result = func(**args)
964
+ response = self._client.chat.completions.create(
965
+ model=self._openai_model,
966
+ messages=[
967
+ {
968
+ "role": "system",
969
+ "content": self._instructions,
970
+ },
971
+ {
972
+ "role": "user",
973
+ "content": result,
974
+ },
975
+ ],
976
+ tools=self._tools,
977
+ stream=True,
978
+ )
979
+ for chunk in response:
980
+ delta = chunk.choices[0].delta
981
+
982
+ if delta.content is not None:
983
+ await self._accumulated_value_queue.put(delta.content)
984
+ # Remove the cached tool call info so it doesn't block future calls
985
+ del final_tool_calls[index]
986
+ except json.JSONDecodeError:
987
+ # If the accumulated arguments aren't valid yet, wait for more chunks.
988
+ continue
989
+
990
+ # Process regular response content
991
+ if delta.content is not None:
992
+ await self._accumulated_value_queue.put(delta.content)
993
+
994
+ # Start the stream processor as a background task
1141
995
  asyncio.create_task(stream_processor())
1142
996
 
1143
- # Collect the full response
1144
- full_response = ""
997
+ # Yield values from the queue as they become available.
1145
998
  while True:
1146
999
  try:
1147
- value = await asyncio.wait_for(
1148
- self._accumulated_value_queue.get(), timeout=0.1
1149
- )
1000
+ value = await asyncio.wait_for(self._accumulated_value_queue.get(), timeout=0.1)
1150
1001
  if value is not None:
1151
- full_response += value
1002
+ final_response += value
1003
+ yield value
1152
1004
  except asyncio.TimeoutError:
1005
+ # Break only if the queue is empty (assuming stream ended)
1153
1006
  if self._accumulated_value_queue.empty():
1154
1007
  break
1155
1008
 
1009
+ # Save the conversation to the database and Zep memory (if configured)
1156
1010
  metadata = {
1157
1011
  "user_id": user_id,
1158
1012
  "message": transcript,
1159
- "response": full_response,
1013
+ "response": final_response,
1160
1014
  "timestamp": datetime.datetime.now(datetime.timezone.utc),
1161
1015
  }
1162
-
1163
1016
  self._database.save_message(user_id, metadata)
1164
-
1165
1017
  if self._zep:
1166
1018
  messages = [
1167
1019
  Message(
@@ -1172,7 +1024,7 @@ class AI:
1172
1024
  Message(
1173
1025
  role="assistant",
1174
1026
  role_type="assistant",
1175
- content=full_response,
1027
+ content=final_response,
1176
1028
  ),
1177
1029
  ]
1178
1030
  await self._zep.memory.add(session_id=user_id, messages=messages)
@@ -1181,32 +1033,12 @@ class AI:
1181
1033
  with self._client.audio.speech.with_streaming_response.create(
1182
1034
  model="tts-1",
1183
1035
  voice=voice,
1184
- input=full_response,
1036
+ input=final_response,
1185
1037
  response_format=response_format,
1186
1038
  ) as response:
1187
1039
  for chunk in response.iter_bytes(1024):
1188
1040
  yield chunk
1189
1041
 
1190
- def _handle_requires_action(self, data, run_id):
1191
- tool_outputs = []
1192
-
1193
- for tool in data.required_action.submit_tool_outputs.tool_calls:
1194
- if tool.function.name in self._tool_handlers:
1195
- handler = self._tool_handlers[tool.function.name]
1196
- inputs = json.loads(tool.function.arguments)
1197
- output = handler(**inputs)
1198
- tool_outputs.append(
1199
- {"tool_call_id": tool.id, "output": output})
1200
-
1201
- self._submit_tool_outputs(tool_outputs, run_id)
1202
-
1203
- def _submit_tool_outputs(self, tool_outputs, run_id):
1204
- with self._client.beta.threads.runs.submit_tool_outputs_stream(
1205
- thread_id=self._current_thread_id, run_id=run_id, tool_outputs=tool_outputs
1206
- ) as stream:
1207
- for text in stream.text_deltas:
1208
- asyncio.create_task(self._accumulated_value_queue.put(text))
1209
-
1210
1042
  def add_tool(self, func: Callable):
1211
1043
  """Register a custom function as an AI tool using decorator pattern.
1212
1044
 
@@ -1253,7 +1085,8 @@ class AI:
1253
1085
  },
1254
1086
  }
1255
1087
  self._tools.append(tool_config)
1256
- self._tool_handlers[func.__name__] = func
1088
+ # Attach the function to the instance so getattr can find it later.
1089
+ setattr(self, func.__name__, func)
1257
1090
  return func
1258
1091
 
1259
1092
 
File without changes