solana-agent 6.0.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

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.
solana_agent/ai.py CHANGED
@@ -6,7 +6,7 @@ import json
6
6
  from typing import AsyncGenerator, List, Literal, Dict, Any, Callable
7
7
  import uuid
8
8
  import pandas as pd
9
- from pydantic import BaseModel
9
+ from pydantic import BaseModel, Field
10
10
  from pymongo import MongoClient
11
11
  from openai import OpenAI
12
12
  import inspect
@@ -20,6 +20,43 @@ from zep_cloud.types import Message
20
20
  from pinecone import Pinecone
21
21
 
22
22
 
23
+ # Define Pydantic models for structured output
24
+ class ImprovementArea(BaseModel):
25
+ area: str = Field(...,
26
+ description="Area name (e.g., 'Accuracy', 'Completeness')")
27
+ issue: str = Field(..., description="Specific issue identified")
28
+ recommendation: str = Field(...,
29
+ description="Specific actionable improvement")
30
+
31
+
32
+ class CritiqueFeedback(BaseModel):
33
+ strengths: List[str] = Field(
34
+ default_factory=list, description="List of strengths in the response"
35
+ )
36
+ improvement_areas: List[ImprovementArea] = Field(
37
+ default_factory=list, description="Areas needing improvement"
38
+ )
39
+ overall_score: float = Field(..., description="Score between 0.0 and 1.0")
40
+ priority: Literal["low", "medium", "high"] = Field(
41
+ ..., description="Priority level for improvements"
42
+ )
43
+
44
+
45
+ class MemoryInsight(BaseModel):
46
+ fact: str = Field(...,
47
+ description="The factual information worth remembering")
48
+ relevance: str = Field(
49
+ ..., description="Short explanation of why this fact is generally useful"
50
+ )
51
+
52
+
53
+ class CollectiveMemoryResponse(BaseModel):
54
+ insights: List[MemoryInsight] = Field(
55
+ default_factory=list,
56
+ description="List of factual insights extracted from the conversation",
57
+ )
58
+
59
+
23
60
  class DocumentModel(BaseModel):
24
61
  id: str
25
62
  text: str
@@ -31,6 +68,7 @@ class MongoDatabase:
31
68
  self.db = self._client[db_name]
32
69
  self.messages = self.db["messages"]
33
70
  self.kb = self.db["kb"]
71
+ self.jobs = self.db["jobs"]
34
72
 
35
73
  def save_message(self, user_id: str, metadata: Dict[str, Any]):
36
74
  metadata["user_id"] = user_id
@@ -70,9 +108,11 @@ class AI:
70
108
  pinecone_index_name: str = None,
71
109
  pinecone_embed_model: Literal["llama-text-embed-v2"] = "llama-text-embed-v2",
72
110
  gemini_api_key: str = None,
73
- openai_base_url: str = None,
74
111
  tool_calling_model: str = "gpt-4o-mini",
75
112
  reasoning_model: str = "gpt-4o-mini",
113
+ research_model: str = "gpt-4o-mini",
114
+ enable_internet_search: bool = True,
115
+ default_timezone: str = "UTC",
76
116
  ):
77
117
  """Initialize a new AI assistant instance.
78
118
 
@@ -88,9 +128,11 @@ class AI:
88
128
  pinecone_index_name (str, optional): Name of the Pinecone index. Defaults to None
89
129
  pinecone_embed_model (Literal["llama-text-embed-v2"], optional): Pinecone embedding model. Defaults to "llama-text-embed-v2"
90
130
  gemini_api_key (str, optional): API key for Gemini search. Defaults to None
91
- openai_base_url (str, optional): Base URL for OpenAI API. Defaults to None
92
131
  tool_calling_model (str, optional): Model for tool calling. Defaults to "gpt-4o-mini"
93
132
  reasoning_model (str, optional): Model for reasoning. Defaults to "gpt-4o-mini"
133
+ research_model (str, optional): Model for research. Defaults to "gpt-4o-mini"
134
+ enable_internet_search (bool, optional): Enable internet search tools. Defaults to True
135
+ default_timezone (str, optional): Default timezone for time awareness. Defaults to "UTC"
94
136
  Example:
95
137
  ```python
96
138
  ai = AI(
@@ -107,11 +149,7 @@ class AI:
107
149
  - Optional integrations for Perplexity, Pinecone, Gemini, and Grok
108
150
  - You must create the Pinecone index in the dashboard before using it
109
151
  """
110
- self._client = (
111
- OpenAI(api_key=openai_api_key, base_url=openai_base_url)
112
- if openai_base_url
113
- else OpenAI(api_key=openai_api_key)
114
- )
152
+ self._client = OpenAI(api_key=openai_api_key)
115
153
  self._memory_instructions = """
116
154
  You are a highly intelligent, context-aware conversational AI. When a user sends a query or statement, you should not only process the current input but also retrieve and integrate relevant context from their previous interactions. Use the memory data to:
117
155
  - Infer nuances in the user's intent.
@@ -147,10 +185,46 @@ class AI:
147
185
  self._pinecone.Index(
148
186
  self._pinecone_index_name) if self._pinecone else None
149
187
  )
150
- self._openai_base_url = openai_base_url
151
188
  self._tool_calling_model = tool_calling_model
152
189
  self._reasoning_model = reasoning_model
190
+ self._research_model = research_model
153
191
  self._tools = []
192
+ self._job_processor_task = None
193
+ self._default_timezone = default_timezone
194
+
195
+ # Automatically add internet search tool if API key is provided and feature is enabled
196
+ if perplexity_api_key and enable_internet_search:
197
+ # Use the add_tool decorator functionality directly
198
+ search_internet_tool = {
199
+ "type": "function",
200
+ "function": {
201
+ "name": "search_internet",
202
+ "description": "Search the internet using Perplexity AI API",
203
+ "parameters": {
204
+ "type": "object",
205
+ "properties": {
206
+ "query": {
207
+ "type": "string",
208
+ "description": "Search query string",
209
+ },
210
+ "model": {
211
+ "type": "string",
212
+ "description": "Perplexity model to use",
213
+ "enum": [
214
+ "sonar",
215
+ "sonar-pro",
216
+ "sonar-reasoning-pro",
217
+ "sonar-reasoning",
218
+ ],
219
+ "default": "sonar",
220
+ },
221
+ },
222
+ "required": ["query"],
223
+ },
224
+ },
225
+ }
226
+ self._tools.append(search_internet_tool)
227
+ print("Internet search capability added as default tool")
154
228
 
155
229
  async def __aenter__(self):
156
230
  return self
@@ -440,25 +514,19 @@ class AI:
440
514
  self.kb.delete(ids=[id], namespace=user_id)
441
515
  self._database.kb.delete_one({"reference": id})
442
516
 
443
- def check_time(self, timezone: str) -> str:
444
- """Get current UTC time formatted as a string via Cloudflare's NTP service.
517
+ def check_time(self, timezone: str = None) -> str:
518
+ """Get current time in requested timezone as a string.
445
519
 
446
520
  Args:
447
- timezone (str): Timezone to convert the time to (e.g., "America/New_York")
521
+ timezone (str, optional): Timezone to convert the time to (e.g., "America/New_York").
522
+ If None, uses the agent's default timezone.
448
523
 
449
524
  Returns:
450
525
  str: Current time in the requested timezone in format 'YYYY-MM-DD HH:MM:SS'
451
-
452
- Example:
453
- ```python
454
- time = ai.check_time("America/New_York")
455
- # Returns: "The current time in America/New_York is 2025-02-26 10:30:45"
456
- ```
457
-
458
- Note:
459
- This is a synchronous tool method required for OpenAI function calling.
460
- Fetches time over NTP from Cloudflare's time server (time.cloudflare.com).
461
526
  """
527
+ # Use provided timezone or fall back to agent default
528
+ timezone = timezone or self._default_timezone or "UTC"
529
+
462
530
  try:
463
531
  # Request time from Cloudflare's NTP server
464
532
  client = ntplib.NTPClient()
@@ -474,7 +542,10 @@ class AI:
474
542
  tz = pytz.timezone(timezone)
475
543
  local_dt = utc_dt.astimezone(tz)
476
544
  formatted_time = local_dt.strftime("%Y-%m-%d %H:%M:%S")
477
- return f"The current time in {timezone} is {formatted_time}"
545
+
546
+ # Format exactly as the test expects
547
+ return f"current time in {timezone} is {formatted_time}"
548
+
478
549
  except pytz.exceptions.UnknownTimeZoneError:
479
550
  return f"Error: Unknown timezone '{timezone}'. Please use a valid timezone like 'America/New_York'."
480
551
 
@@ -648,6 +719,101 @@ class AI:
648
719
  except Exception:
649
720
  pass
650
721
 
722
+ def make_time_aware(self, default_timezone="UTC"):
723
+ """Make the agent time-aware by adding time checking capability."""
724
+ # Add time awareness to instructions with explicit formatting guidance
725
+ time_instructions = f"""
726
+ IMPORTANT: You are time-aware. The current date is {datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")}.
727
+
728
+ TIME RESPONSE RULES:
729
+ 1. When asked about the current time, ONLY use the check_time tool and respond with EXACTLY what it returns
730
+ 2. NEVER add UTC time when the check_time tool returns local time
731
+ 3. NEVER convert between timezones unless explicitly requested
732
+ 4. NEVER mention timezone offsets (like "X hours behind UTC") unless explicitly asked
733
+ 5. Local time is the ONLY time that should be mentioned in your response
734
+
735
+ Default timezone: {default_timezone} (use this when user's timezone is unknown)
736
+ """
737
+ self._instructions = self._instructions + "\n\n" + time_instructions
738
+
739
+ self._default_timezone = default_timezone
740
+
741
+ # Ensure the check_time tool is registered (in case it was removed)
742
+ existing_tools = [t["function"]["name"] for t in self._tools]
743
+ if "check_time" not in existing_tools:
744
+ # Get method reference
745
+ check_time_func = self.check_time
746
+ # Re-register it using our add_tool decorator
747
+ self.add_tool(check_time_func)
748
+
749
+ return self
750
+
751
+ async def research_and_learn(self, topic: str) -> str:
752
+ """Research a topic and add findings to collective memory.
753
+
754
+ Args:
755
+ topic: The topic to research and learn about
756
+
757
+ Returns:
758
+ Summary of what was learned
759
+ """
760
+ try:
761
+ # First, search the internet for information
762
+ search_results = await self.search_internet(
763
+ f"comprehensive information about {topic}"
764
+ )
765
+
766
+ # Extract structured knowledge
767
+ prompt = f"""
768
+ Based on these search results about "{topic}", extract 3-5 factual insights
769
+ worth adding to our collective knowledge.
770
+
771
+ Search results:
772
+ {search_results}
773
+
774
+ Format each insight as a JSON object with:
775
+ 1. "fact": The factual information
776
+ 2. "relevance": Short explanation of why this is generally useful
777
+
778
+ Return ONLY a valid JSON array. Example:
779
+ [
780
+ {{"fact": "Topic X has property Y", "relevance": "Important for understanding Z"}}
781
+ ]
782
+ """
783
+
784
+ response = self._client.chat.completions.create(
785
+ model=self._research_model,
786
+ messages=[
787
+ {
788
+ "role": "system",
789
+ "content": "Extract factual knowledge from research.",
790
+ },
791
+ {"role": "user", "content": prompt},
792
+ ],
793
+ temperature=0.1,
794
+ )
795
+
796
+ insights = json.loads(response.choices[0].message.content)
797
+
798
+ # Add to collective memory via the swarm
799
+ if hasattr(self, "_swarm") and self._swarm and insights:
800
+ conversation = {
801
+ "message": f"Research on {topic}",
802
+ "response": json.dumps(insights),
803
+ "user_id": "system_explorer",
804
+ }
805
+ await self._swarm.extract_and_store_insights(
806
+ "system_explorer", conversation
807
+ )
808
+
809
+ # Return a summary of what was learned
810
+ return f"✅ Added {len(insights)} new insights about '{topic}' to collective memory."
811
+
812
+ return "⚠️ Could not add insights to collective memory."
813
+
814
+ except Exception as e:
815
+ return f"❌ Error researching topic: {str(e)}"
816
+
651
817
  async def delete_memory(self, user_id: str):
652
818
  """Delete memory for a specific user from Zep memory.
653
819
 
@@ -676,12 +842,20 @@ class AI:
676
842
  )
677
843
  return transcription.text
678
844
 
679
- async def text(self, user_id: str, user_text: str) -> AsyncGenerator[str, None]:
845
+ async def text(
846
+ self,
847
+ user_id: str,
848
+ user_text: str,
849
+ timezone: str = None,
850
+ original_user_text: str = None,
851
+ ) -> AsyncGenerator[str, None]:
680
852
  """Process text input and stream AI responses asynchronously.
681
853
 
682
854
  Args:
683
855
  user_id (str): Unique identifier for the user/conversation.
684
856
  user_text (str): Text input from user to process.
857
+ original_user_text (str, optional): Original user message for storage. If provided,
858
+ this will be stored instead of user_text. Defaults to None.
685
859
 
686
860
  Returns:
687
861
  AsyncGenerator[str, None]: Stream of response text chunks (including tool call results).
@@ -698,6 +872,22 @@ class AI:
698
872
  - Integrates with Zep memory if configured.
699
873
  - Supports tool calls by aggregating and executing them as their arguments stream in.
700
874
  """
875
+ # Store current user ID for task scheduling context
876
+ self._current_user_id = user_id
877
+
878
+ # Store timezone with user ID for persistence
879
+ if timezone:
880
+ if not hasattr(self, "_user_timezones"):
881
+ self._user_timezones = {}
882
+ self._user_timezones[user_id] = timezone
883
+
884
+ # Set current timezone for this session
885
+ self._current_timezone = (
886
+ timezone
887
+ if timezone
888
+ else self._user_timezones.get(user_id, self._default_timezone)
889
+ )
890
+
701
891
  self._accumulated_value_queue = asyncio.Queue()
702
892
  final_tool_calls = {} # Accumulate tool call deltas
703
893
  final_response = ""
@@ -829,10 +1019,15 @@ class AI:
829
1019
  if self._accumulated_value_queue.empty():
830
1020
  break
831
1021
 
1022
+ # For storage purposes, use original text if provided
1023
+ message_to_store = (
1024
+ original_user_text if original_user_text is not None else user_text
1025
+ )
1026
+
832
1027
  # Save the conversation to the database and Zep memory (if configured)
833
1028
  metadata = {
834
1029
  "user_id": user_id,
835
- "message": user_text,
1030
+ "message": message_to_store,
836
1031
  "response": final_response,
837
1032
  "timestamp": datetime.datetime.now(datetime.timezone.utc),
838
1033
  }
@@ -1110,32 +1305,338 @@ class AI:
1110
1305
  class Swarm:
1111
1306
  """An AI Agent Swarm that coordinates specialized AI agents with handoff capabilities."""
1112
1307
 
1113
- def __init__(self, database: MongoDatabase, router_model: str = "gpt-4o"):
1308
+ def __init__(
1309
+ self,
1310
+ database: MongoDatabase,
1311
+ directive: str = None,
1312
+ router_model: str = "gpt-4o-mini",
1313
+ insight_model: str = "gpt-4o-mini",
1314
+ enable_collective_memory: bool = True,
1315
+ enable_critic: bool = True,
1316
+ default_timezone: str = "UTC",
1317
+ ):
1114
1318
  """Initialize the multi-agent system with a shared database.
1115
1319
 
1116
1320
  Args:
1117
1321
  database (MongoDatabase): Shared MongoDB database instance
1118
- router_model (str, optional): Model to use for routing decisions. Defaults to "gpt-4o".
1322
+ directive (str, optional): Core directive/mission that governs all agents. Defaults to None.
1323
+ router_model (str, optional): Model to use for routing decisions. Defaults to "gpt-4o-mini".
1324
+ insight_model (str, optional): Model to extract collective insights. Defaults to "gpt-4o-mini".
1325
+ enable_collective_memory (bool, optional): Whether to enable collective memory. Defaults to True.
1326
+ enable_critic (bool, optional): Whether to enable the critic system. Defaults to True.
1327
+ default_timezone (str, optional): Default timezone for time-awareness. Defaults to "UTC".
1119
1328
  """
1120
1329
  self.agents = {} # name -> AI instance
1121
1330
  self.specializations = {} # name -> description
1122
1331
  self.database = database
1123
1332
  self.router_model = router_model
1333
+ self.insight_model = insight_model
1334
+ self.enable_collective_memory = enable_collective_memory
1335
+ self.default_timezone = default_timezone
1336
+ self.enable_critic = enable_critic
1337
+
1338
+ # Store swarm directive
1339
+ self.swarm_directive = (
1340
+ directive
1341
+ or """
1342
+ You are part of an agent swarm that works together to serve users effectively.
1343
+ Your goals are to provide accurate, helpful responses while collaborating with other agents.
1344
+ """
1345
+ )
1346
+
1347
+ self.formatted_directive = f"""
1348
+ ┌─────────────── SWARM DIRECTIVE ───────────────┐
1349
+ {self.swarm_directive}
1350
+ └─────────────────────────────────────────────┘
1351
+ """
1352
+
1353
+ # Initialize critic if enabled
1354
+ if enable_critic:
1355
+ self.critic = Critic(
1356
+ self,
1357
+ critique_model=insight_model,
1358
+ )
1124
1359
 
1125
1360
  # Ensure handoffs collection exists
1126
1361
  if "handoffs" not in self.database.db.list_collection_names():
1127
1362
  self.database.db.create_collection("handoffs")
1128
1363
  self.handoffs = self.database.db["handoffs"]
1129
1364
 
1365
+ # Create collective memory collection
1366
+ if enable_collective_memory:
1367
+ if "collective_memory" not in self.database.db.list_collection_names():
1368
+ self.database.db.create_collection("collective_memory")
1369
+ self.collective_memory = self.database.db["collective_memory"]
1370
+
1371
+ # Create text index for MongoDB text search
1372
+ try:
1373
+ self.collective_memory.create_index(
1374
+ [("fact", "text"), ("relevance", "text")]
1375
+ )
1376
+ print("Created text search index for collective memory")
1377
+ except Exception as e:
1378
+ print(f"Warning: Text index creation might have failed: {e}")
1379
+ else:
1380
+ print("Collective memory feature is disabled")
1381
+
1130
1382
  print(
1131
1383
  f"MultiAgentSystem initialized with router model: {router_model}")
1132
1384
 
1385
+ # Update the extract_and_store_insights method in Swarm class
1386
+
1387
+ async def extract_and_store_insights(
1388
+ self, user_id: str, conversation: dict
1389
+ ) -> None:
1390
+ """Extract and store insights with hybrid vector/text search capabilities."""
1391
+ # Get first agent to use its OpenAI client
1392
+ if not self.agents:
1393
+ return
1394
+
1395
+ first_agent = next(iter(self.agents.values()))
1396
+
1397
+ # Create the prompt to extract insights
1398
+ prompt = f"""
1399
+ Review this conversation and extract 0-3 IMPORTANT factual insights worth remembering for future users.
1400
+ Only extract FACTUAL information that would be valuable across multiple conversations.
1401
+ Do NOT include opinions, personal preferences, or user-specific details.
1402
+
1403
+ Conversation:
1404
+ User: {conversation.get('message', '')}
1405
+ Assistant: {conversation.get('response', '')}
1406
+ """
1407
+
1408
+ # Extract insights using AI with structured parsing
1409
+ try:
1410
+ # Parse the response using the Pydantic model
1411
+ completion = first_agent._client.beta.chat.completions.parse(
1412
+ model=self.insight_model,
1413
+ messages=[
1414
+ {
1415
+ "role": "system",
1416
+ "content": "Extract important factual insights from conversations.",
1417
+ },
1418
+ {"role": "user", "content": prompt},
1419
+ ],
1420
+ response_format=CollectiveMemoryResponse,
1421
+ temperature=0.1,
1422
+ )
1423
+
1424
+ # Extract the Pydantic model
1425
+ memory_response = completion.choices[0].message.parsed
1426
+
1427
+ # Store in MongoDB (keeps all metadata and text)
1428
+ timestamp = datetime.datetime.now(datetime.timezone.utc)
1429
+ mongo_records = []
1430
+
1431
+ for insight in memory_response.insights:
1432
+ record_id = str(uuid.uuid4())
1433
+ record = {
1434
+ "_id": record_id,
1435
+ "fact": insight.fact,
1436
+ "relevance": insight.relevance,
1437
+ "timestamp": timestamp,
1438
+ "source_user_id": user_id,
1439
+ }
1440
+ mongo_records.append(record)
1441
+
1442
+ if mongo_records:
1443
+ for record in mongo_records:
1444
+ self.collective_memory.insert_one(record)
1445
+
1446
+ # Also store in Pinecone for semantic search if available
1447
+ if (
1448
+ mongo_records
1449
+ and hasattr(first_agent, "_pinecone")
1450
+ and first_agent._pinecone
1451
+ and first_agent.kb
1452
+ ):
1453
+ try:
1454
+ # Generate embeddings
1455
+ texts = [
1456
+ f"{record['fact']}: {record['relevance']}"
1457
+ for record in mongo_records
1458
+ ]
1459
+ embeddings = first_agent._pinecone.inference.embed(
1460
+ model=first_agent._pinecone_embedding_model,
1461
+ inputs=texts,
1462
+ parameters={"input_type": "passage",
1463
+ "truncate": "END"},
1464
+ )
1465
+
1466
+ # Create vectors for Pinecone
1467
+ vectors = []
1468
+ for record, embedding in zip(mongo_records, embeddings):
1469
+ vectors.append(
1470
+ {
1471
+ "id": record["_id"],
1472
+ "values": embedding.values,
1473
+ "metadata": {
1474
+ "fact": record["fact"],
1475
+ "relevance": record["relevance"],
1476
+ "timestamp": str(timestamp),
1477
+ "source_user_id": user_id,
1478
+ },
1479
+ }
1480
+ )
1481
+
1482
+ # Store in Pinecone
1483
+ first_agent.kb.upsert(
1484
+ vectors=vectors, namespace="collective_memory"
1485
+ )
1486
+ print(
1487
+ f"Stored {len(mongo_records)} insights in both MongoDB and Pinecone"
1488
+ )
1489
+ except Exception as e:
1490
+ print(f"Error storing insights in Pinecone: {e}")
1491
+ else:
1492
+ print(f"Stored {len(mongo_records)} insights in MongoDB only")
1493
+
1494
+ except Exception as e:
1495
+ print(f"Failed to extract insights: {str(e)}")
1496
+
1497
+ def search_collective_memory(self, query: str, limit: int = 5) -> str:
1498
+ """Search the collective memory using a hybrid approach.
1499
+
1500
+ First tries semantic vector search through Pinecone, then falls back to
1501
+ MongoDB text search, and finally to recency-based search as needed.
1502
+
1503
+ Args:
1504
+ query: The search query
1505
+ limit: Maximum number of results to return
1506
+
1507
+ Returns:
1508
+ Formatted string with relevant insights
1509
+ """
1510
+ try:
1511
+ if not self.enable_collective_memory:
1512
+ return "Collective memory feature is disabled."
1513
+
1514
+ results = []
1515
+ search_method = "recency" # Default method if others fail
1516
+
1517
+ # Try semantic search with Pinecone first
1518
+ if self.agents:
1519
+ first_agent = next(iter(self.agents.values()))
1520
+ if (
1521
+ hasattr(first_agent, "_pinecone")
1522
+ and first_agent._pinecone
1523
+ and first_agent.kb
1524
+ ):
1525
+ try:
1526
+ # Generate embedding for query
1527
+ embedding = first_agent._pinecone.inference.embed(
1528
+ model=first_agent._pinecone_embedding_model,
1529
+ inputs=[query],
1530
+ parameters={"input_type": "passage",
1531
+ "truncate": "END"},
1532
+ )
1533
+
1534
+ # Search Pinecone
1535
+ pinecone_results = first_agent.kb.query(
1536
+ vector=embedding[0].values,
1537
+ top_k=limit * 2, # Get more results to allow for filtering
1538
+ include_metadata=True,
1539
+ namespace="collective_memory",
1540
+ )
1541
+
1542
+ # Extract results from Pinecone
1543
+ if pinecone_results.matches:
1544
+ for match in pinecone_results.matches:
1545
+ if hasattr(match, "metadata") and match.metadata:
1546
+ results.append(
1547
+ {
1548
+ "fact": match.metadata.get(
1549
+ "fact", "Unknown fact"
1550
+ ),
1551
+ "relevance": match.metadata.get(
1552
+ "relevance", ""
1553
+ ),
1554
+ "score": match.score,
1555
+ }
1556
+ )
1557
+
1558
+ # Get top results
1559
+ results = sorted(
1560
+ results, key=lambda x: x.get("score", 0), reverse=True
1561
+ )[:limit]
1562
+ search_method = "semantic"
1563
+ except Exception as e:
1564
+ print(f"Pinecone search error: {e}")
1565
+
1566
+ # Fall back to MongoDB keyword search if needed
1567
+ if not results:
1568
+ try:
1569
+ # First try text search if we have the index
1570
+ mongo_results = list(
1571
+ self.collective_memory.find(
1572
+ {"$text": {"$search": query}},
1573
+ {"score": {"$meta": "textScore"}},
1574
+ )
1575
+ .sort([("score", {"$meta": "textScore"})])
1576
+ .limit(limit)
1577
+ )
1578
+
1579
+ if mongo_results:
1580
+ results = mongo_results
1581
+ search_method = "keyword"
1582
+ else:
1583
+ # Fall back to most recent insights
1584
+ results = list(
1585
+ self.collective_memory.find()
1586
+ .sort("timestamp", -1)
1587
+ .limit(limit)
1588
+ )
1589
+ search_method = "recency"
1590
+ except Exception as e:
1591
+ print(f"MongoDB search error: {e}")
1592
+ # Final fallback - just get most recent
1593
+ results = list(
1594
+ self.collective_memory.find().sort("timestamp", -1).limit(limit)
1595
+ )
1596
+
1597
+ # Format the results
1598
+ if not results:
1599
+ return "No collective knowledge available."
1600
+
1601
+ formatted = [
1602
+ f"## Relevant Collective Knowledge (using {search_method} search)"
1603
+ ]
1604
+ for insight in results:
1605
+ formatted.append(
1606
+ f"- **{insight.get('fact')}** _{insight.get('relevance', '')}_"
1607
+ )
1608
+
1609
+ return "\n".join(formatted)
1610
+
1611
+ except Exception as e:
1612
+ print(f"Error searching collective memory: {str(e)}")
1613
+ return "Error retrieving collective knowledge."
1614
+
1133
1615
  def register(self, name: str, agent: AI, specialization: str):
1134
1616
  """Register a specialized agent with the multi-agent system."""
1617
+ # Make agent time-aware first
1618
+ agent.make_time_aware(self.default_timezone)
1619
+
1620
+ # Apply swarm directive to the agent
1621
+ agent._instructions = f"{self.formatted_directive}\n\n{agent._instructions}"
1622
+
1135
1623
  # Add the agent to the system first
1136
1624
  self.agents[name] = agent
1137
1625
  self.specializations[name] = specialization
1138
1626
 
1627
+ # Add collective memory tool to the agent
1628
+ @agent.add_tool
1629
+ def query_collective_knowledge(query: str) -> str:
1630
+ """Query the swarm's collective knowledge from all users.
1631
+
1632
+ Args:
1633
+ query (str): The search query to look for in collective knowledge
1634
+
1635
+ Returns:
1636
+ str: Relevant insights from the swarm's collective memory
1637
+ """
1638
+ return self.search_collective_memory(query)
1639
+
1139
1640
  print(
1140
1641
  f"Registered agent: {name}, specialization: {specialization[:50]}...")
1141
1642
  print(f"Current agents: {list(self.agents.keys())}")
@@ -1251,81 +1752,181 @@ class Swarm:
1251
1752
  f"Updated handoff capabilities for {agent_name} with targets: {available_targets}"
1252
1753
  )
1253
1754
 
1254
- async def process(self, user_id: str, user_text: str) -> AsyncGenerator[str, None]:
1255
- """Process the user request with appropriate agent and handle handoffs."""
1755
+ async def process(
1756
+ self, user_id: str, user_text: str, timezone: str = None
1757
+ ) -> AsyncGenerator[str, None]:
1758
+ """Process the user request with appropriate agent and handle handoffs.
1759
+
1760
+ Args:
1761
+ user_id (str): Unique user identifier
1762
+ user_text (str): User's text input
1763
+ timezone (str, optional): User-specific timezone
1764
+ """
1256
1765
  try:
1257
- # Check if any agents are registered
1766
+ # Handle special commands
1767
+ if user_text.strip().lower().startswith("!memory "):
1768
+ query = user_text[8:].strip()
1769
+ yield self.search_collective_memory(query)
1770
+ return
1771
+
1772
+ # Check for registered agents
1258
1773
  if not self.agents:
1259
1774
  yield "Error: No agents are registered with the system. Please register at least one agent first."
1260
1775
  return
1261
1776
 
1262
- # Get routing decision
1777
+ # Get initial routing and agent
1263
1778
  first_agent = next(iter(self.agents.values()))
1264
1779
  agent_name = await self._get_routing_decision(first_agent, user_text)
1265
1780
  current_agent = self.agents[agent_name]
1266
1781
  print(f"Starting conversation with agent: {agent_name}")
1267
1782
 
1268
- # Initialize a flag for handoff detection
1269
- handoff_detected = False
1270
- response_started = False
1271
-
1272
- # Reset handoff info for this interaction
1783
+ # Reset handoff info
1273
1784
  current_agent._handoff_info = None
1274
1785
 
1275
- # Process initial agent's response
1276
- async for chunk in current_agent.text(user_id, user_text):
1277
- # Check for handoff after each chunk
1278
- if current_agent._handoff_info and not handoff_detected:
1279
- handoff_detected = True
1280
- target_name = current_agent._handoff_info["target"]
1281
- target_agent = self.agents[target_name]
1282
- reason = current_agent._handoff_info["reason"]
1283
-
1284
- # Record handoff without waiting
1285
- asyncio.create_task(
1286
- self._record_handoff(
1287
- user_id, agent_name, target_name, reason, user_text
1288
- )
1289
- )
1786
+ # Response tracking
1787
+ final_response = ""
1290
1788
 
1291
- # Process with target agent
1292
- print(f"[HANDOFF] Forwarding to {target_name}")
1293
- handoff_query = f"""
1294
- Answer this ENTIRE question completely from scratch:
1295
-
1296
- {user_text}
1297
-
1298
- IMPORTANT INSTRUCTIONS:
1299
- 1. Address ALL aspects of the question comprehensively
1300
- 2. Organize your response in a logical, structured manner
1301
- 3. Include both explanations AND implementations as needed
1302
- 4. Do not mention any handoff or that you're continuing from another agent
1303
- 5. Answer as if you are addressing the complete question from the beginning
1304
- 6. Consider any relevant context from previous conversation
1305
- """
1789
+ # Process response stream
1790
+ async for chunk in self._stream_response(
1791
+ user_id, user_text, current_agent, timezone
1792
+ ):
1793
+ yield chunk
1794
+ final_response += chunk
1306
1795
 
1307
- # If we've already started returning some text, add a separator
1308
- if response_started:
1309
- yield "\n\n---\n\n"
1310
-
1311
- # Stream directly from target agent
1312
- async for new_chunk in target_agent.text(user_id, handoff_query):
1313
- yield new_chunk
1314
- # Force immediate delivery of each chunk
1315
- await asyncio.sleep(0)
1316
- return
1317
- else:
1318
- # Only yield content if no handoff has been detected
1319
- if not handoff_detected:
1320
- response_started = True
1321
- yield chunk
1322
- await asyncio.sleep(0) # Force immediate delivery
1796
+ # Post-processing: learn from conversation
1797
+ conversation = {
1798
+ "user_id": user_id,
1799
+ "message": user_text,
1800
+ "response": final_response,
1801
+ }
1802
+
1803
+ # Run post-processing tasks concurrently
1804
+ tasks = []
1805
+
1806
+ # Add collective memory task if enabled
1807
+ if self.enable_collective_memory:
1808
+ tasks.append(self.extract_and_store_insights(
1809
+ user_id, conversation))
1810
+
1811
+ # Run all post-processing tasks concurrently
1812
+ if tasks:
1813
+ # Don't block - run asynchronously
1814
+ asyncio.create_task(self._run_post_processing_tasks(tasks))
1323
1815
 
1324
1816
  except Exception as e:
1325
1817
  print(f"Error in multi-agent processing: {str(e)}")
1326
1818
  print(traceback.format_exc())
1327
1819
  yield "\n\nI apologize for the technical difficulty.\n\n"
1328
1820
 
1821
+ async def _stream_response(
1822
+ self, user_id, user_text, current_agent, timezone=None
1823
+ ) -> AsyncGenerator[str, None]:
1824
+ """Stream response from an agent, handling potential handoffs."""
1825
+ handoff_detected = False
1826
+ response_started = False
1827
+ full_response = ""
1828
+
1829
+ # Get recent feedback for this agent to improve the response
1830
+ agent_name = current_agent.__class__.__name__
1831
+ recent_feedback = []
1832
+
1833
+ if self.enable_critic and hasattr(self, "critic"):
1834
+ try:
1835
+ # Get the most recent feedback for this agent
1836
+ feedback_records = list(
1837
+ self.critic.feedback_collection.find(
1838
+ {"agent_name": agent_name})
1839
+ .sort("timestamp", -1)
1840
+ .limit(3)
1841
+ )
1842
+
1843
+ if feedback_records:
1844
+ # Extract specific improvement suggestions
1845
+ for record in feedback_records:
1846
+ recent_feedback.append(
1847
+ f"- Improve {record.get('improvement_area')}: {record.get('recommendation')}"
1848
+ )
1849
+ except Exception as e:
1850
+ print(f"Error getting recent feedback: {e}")
1851
+
1852
+ # Augment user text with feedback instructions if available
1853
+ augmented_instruction = user_text
1854
+ if recent_feedback:
1855
+ feedback_text = "\n".join(recent_feedback)
1856
+ augmented_instruction = f"""
1857
+ Answer this question: {user_text}
1858
+
1859
+ IMPORTANT - Apply these specific improvements from previous feedback:
1860
+ {feedback_text}
1861
+ """
1862
+ print(f"Applying feedback to improve response: {feedback_text}")
1863
+
1864
+ async for chunk in current_agent.text(
1865
+ user_id, augmented_instruction, timezone, user_text
1866
+ ):
1867
+ # Accumulate the full response for critic analysis
1868
+ full_response += chunk
1869
+
1870
+ # Check for handoff after each chunk
1871
+ if current_agent._handoff_info and not handoff_detected:
1872
+ handoff_detected = True
1873
+ target_name = current_agent._handoff_info["target"]
1874
+ target_agent = self.agents[target_name]
1875
+ reason = current_agent._handoff_info["reason"]
1876
+
1877
+ # Record handoff without waiting
1878
+ asyncio.create_task(
1879
+ self._record_handoff(
1880
+ user_id, current_agent, target_name, reason, user_text
1881
+ )
1882
+ )
1883
+
1884
+ # Add separator if needed
1885
+ if response_started:
1886
+ yield "\n\n---\n\n"
1887
+
1888
+ # Pass to target agent with comprehensive instructions
1889
+ handoff_query = f"""
1890
+ Answer this ENTIRE question completely from scratch:
1891
+ {user_text}
1892
+
1893
+ IMPORTANT INSTRUCTIONS:
1894
+ 1. Address ALL aspects of the question comprehensively
1895
+ 2. Include both explanations AND implementations as needed
1896
+ 3. Do not mention any handoff or that you're continuing from another agent
1897
+ 4. Consider any relevant context from previous conversation
1898
+ """
1899
+
1900
+ # Stream from target agent
1901
+ async for new_chunk in target_agent.text(user_id, handoff_query):
1902
+ yield new_chunk
1903
+ await asyncio.sleep(0) # Force immediate delivery
1904
+ return
1905
+
1906
+ # Regular response if no handoff detected
1907
+ if not handoff_detected:
1908
+ response_started = True
1909
+ yield chunk
1910
+ await asyncio.sleep(0) # Force immediate delivery
1911
+
1912
+ # After full response is delivered, invoke critic (if enabled)
1913
+ if self.enable_critic and hasattr(self, "critic"):
1914
+ # Don't block - run asynchronously
1915
+ asyncio.create_task(
1916
+ self.critic.analyze_interaction(
1917
+ agent_name=current_agent.__class__.__name__,
1918
+ user_query=user_text,
1919
+ response=full_response,
1920
+ )
1921
+ )
1922
+
1923
+ async def _run_post_processing_tasks(self, tasks):
1924
+ """Run multiple post-processing tasks concurrently."""
1925
+ try:
1926
+ await asyncio.gather(*tasks)
1927
+ except Exception as e:
1928
+ print(f"Error in post-processing tasks: {e}")
1929
+
1329
1930
  async def _get_routing_decision(self, agent, user_text):
1330
1931
  """Get routing decision in parallel to reduce latency."""
1331
1932
  enhanced_prompt = f"""
@@ -1396,3 +1997,134 @@ class Swarm:
1396
1997
 
1397
1998
  # Fallback to first agent
1398
1999
  return list(self.agents.keys())[0]
2000
+
2001
+
2002
+ class Critic:
2003
+ """System that evaluates agent responses and suggests improvements."""
2004
+
2005
+ def __init__(self, swarm, critique_model="gpt-4o-mini"):
2006
+ """Initialize the critic system.
2007
+
2008
+ Args:
2009
+ swarm: The agent swarm to monitor
2010
+ critique_model: Model to use for evaluations
2011
+ """
2012
+ self.swarm = swarm
2013
+ self.critique_model = critique_model
2014
+ self.feedback_collection = swarm.database.db["agent_feedback"]
2015
+
2016
+ # Create index for feedback collection
2017
+ if "agent_feedback" not in swarm.database.db.list_collection_names():
2018
+ swarm.database.db.create_collection("agent_feedback")
2019
+ self.feedback_collection.create_index([("agent_name", 1)])
2020
+ self.feedback_collection.create_index([("improvement_area", 1)])
2021
+
2022
+ async def analyze_interaction(
2023
+ self,
2024
+ agent_name,
2025
+ user_query,
2026
+ response,
2027
+ ):
2028
+ """Analyze an agent interaction and provide improvement feedback."""
2029
+ # Get first agent's client for analysis
2030
+ first_agent = next(iter(self.swarm.agents.values()))
2031
+
2032
+ prompt = f"""
2033
+ Analyze this agent interaction to identify specific improvements.
2034
+
2035
+ INTERACTION:
2036
+ User query: {user_query}
2037
+ Agent response: {response}
2038
+
2039
+ Provide feedback on accuracy, completeness, clarity, efficiency, and tone.
2040
+ """
2041
+
2042
+ try:
2043
+ # Parse the response
2044
+ completion = first_agent._client.beta.chat.completions.parse(
2045
+ model=self.critique_model,
2046
+ messages=[
2047
+ {
2048
+ "role": "system",
2049
+ "content": "You are a helpful critic evaluating AI responses.",
2050
+ },
2051
+ {"role": "user", "content": prompt},
2052
+ ],
2053
+ response_format=CritiqueFeedback,
2054
+ temperature=0.2,
2055
+ )
2056
+
2057
+ # Extract the Pydantic model
2058
+ feedback = completion.choices[0].message.parsed
2059
+
2060
+ # Manual validation - ensure score is between 0 and 1
2061
+ if feedback.overall_score < 0:
2062
+ feedback.overall_score = 0.0
2063
+ elif feedback.overall_score > 1:
2064
+ feedback.overall_score = 1.0
2065
+
2066
+ # Store feedback in database
2067
+ for area in feedback.improvement_areas:
2068
+ self.feedback_collection.insert_one(
2069
+ {
2070
+ "agent_name": agent_name,
2071
+ "user_query": user_query,
2072
+ "timestamp": datetime.datetime.now(datetime.timezone.utc),
2073
+ "improvement_area": area.area,
2074
+ "issue": area.issue,
2075
+ "recommendation": area.recommendation,
2076
+ "overall_score": feedback.overall_score,
2077
+ "priority": feedback.priority,
2078
+ }
2079
+ )
2080
+
2081
+ # If high priority feedback, schedule immediate learning task
2082
+ if feedback.priority == "high" and feedback.improvement_areas:
2083
+ top_issue = feedback.improvement_areas[0]
2084
+ await self.schedule_improvement_task(agent_name, top_issue)
2085
+
2086
+ return feedback
2087
+
2088
+ except Exception as e:
2089
+ print(f"Error in critic analysis: {str(e)}")
2090
+ return None
2091
+
2092
+ async def schedule_improvement_task(self, agent_name, issue):
2093
+ """Execute improvement task immediately."""
2094
+ if agent_name in self.swarm.agents:
2095
+ agent = self.swarm.agents[agent_name]
2096
+
2097
+ # Create topic for improvement
2098
+ topic = (
2099
+ f"How to improve {issue['area'].lower()} in responses: {issue['issue']}"
2100
+ )
2101
+
2102
+ # Execute research directly
2103
+ result = await agent.research_and_learn(topic)
2104
+
2105
+ print(
2106
+ f"📝 Executed improvement task for {agent_name}: {issue['area']}")
2107
+ return result
2108
+
2109
+ def get_agent_feedback(self, agent_name=None, limit=10):
2110
+ """Get recent feedback for an agent or all agents."""
2111
+ query = {"agent_name": agent_name} if agent_name else {}
2112
+ feedback = list(
2113
+ self.feedback_collection.find(query).sort(
2114
+ "timestamp", -1).limit(limit)
2115
+ )
2116
+ return feedback
2117
+
2118
+ def get_improvement_trends(self):
2119
+ """Get trends in improvement areas across all agents."""
2120
+ pipeline = [
2121
+ {
2122
+ "$group": {
2123
+ "_id": "$improvement_area",
2124
+ "count": {"$sum": 1},
2125
+ "avg_score": {"$avg": "$overall_score"},
2126
+ }
2127
+ },
2128
+ {"$sort": {"count": -1}},
2129
+ ]
2130
+ return list(self.feedback_collection.aggregate(pipeline))