ambivo-agents 1.0.1__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.
@@ -0,0 +1,1024 @@
1
+ # ambivo_agents/core/base.py - ENHANCED with chat() method
2
+ """
3
+ Enhanced BaseAgent with built-in auto-context session management and simplified chat interface
4
+ """
5
+
6
+ import asyncio
7
+ import uuid
8
+ import time
9
+ import tempfile
10
+ import os
11
+ from pathlib import Path
12
+ from abc import ABC, abstractmethod
13
+ from dataclasses import dataclass, field
14
+ from datetime import datetime, timedelta
15
+ from enum import Enum
16
+ from typing import Dict, List, Any, Optional, Callable, Tuple, Union
17
+ from concurrent.futures import ThreadPoolExecutor
18
+ import logging
19
+
20
+ # Docker imports
21
+ try:
22
+ import docker
23
+
24
+ DOCKER_AVAILABLE = True
25
+ except ImportError:
26
+ DOCKER_AVAILABLE = False
27
+
28
+
29
+ class AgentRole(Enum):
30
+ ASSISTANT = "assistant"
31
+ PROXY = "proxy"
32
+ ANALYST = "analyst"
33
+ RESEARCHER = "researcher"
34
+ COORDINATOR = "coordinator"
35
+ VALIDATOR = "validator"
36
+ CODE_EXECUTOR = "code_executor"
37
+
38
+
39
+ class MessageType(Enum):
40
+ USER_INPUT = "user_input"
41
+ AGENT_RESPONSE = "agent_response"
42
+ SYSTEM_MESSAGE = "system_message"
43
+ TOOL_CALL = "tool_call"
44
+ TOOL_RESPONSE = "tool_response"
45
+ ERROR = "error"
46
+ STATUS_UPDATE = "status_update"
47
+
48
+
49
+ @dataclass
50
+ class AgentMessage:
51
+ id: str
52
+ sender_id: str
53
+ recipient_id: Optional[str]
54
+ content: str
55
+ message_type: MessageType
56
+ metadata: Dict[str, Any] = field(default_factory=dict)
57
+ timestamp: datetime = field(default_factory=datetime.now)
58
+ session_id: Optional[str] = None
59
+ conversation_id: Optional[str] = None
60
+
61
+ def to_dict(self) -> Dict[str, Any]:
62
+ """Convert to dictionary for serialization"""
63
+ return {
64
+ 'id': self.id,
65
+ 'sender_id': self.sender_id,
66
+ 'recipient_id': self.recipient_id,
67
+ 'content': self.content,
68
+ 'message_type': self.message_type.value,
69
+ 'metadata': self.metadata,
70
+ 'timestamp': self.timestamp.isoformat(),
71
+ 'session_id': self.session_id,
72
+ 'conversation_id': self.conversation_id
73
+ }
74
+
75
+ @classmethod
76
+ def from_dict(cls, data: Dict[str, Any]) -> 'AgentMessage':
77
+ """Create from dictionary"""
78
+ return cls(
79
+ id=data['id'],
80
+ sender_id=data['sender_id'],
81
+ recipient_id=data.get('recipient_id'),
82
+ content=data['content'],
83
+ message_type=MessageType(data['message_type']),
84
+ metadata=data.get('metadata', {}),
85
+ timestamp=datetime.fromisoformat(data['timestamp']),
86
+ session_id=data.get('session_id'),
87
+ conversation_id=data.get('conversation_id')
88
+ )
89
+
90
+
91
+ @dataclass
92
+ class AgentTool:
93
+ name: str
94
+ description: str
95
+ function: Callable
96
+ parameters_schema: Dict[str, Any]
97
+ requires_approval: bool = False
98
+ timeout: int = 30
99
+
100
+
101
+ @dataclass
102
+ class ExecutionContext:
103
+ session_id: str
104
+ conversation_id: str
105
+ user_id: str
106
+ tenant_id: str
107
+ metadata: Dict[str, Any] = field(default_factory=dict)
108
+
109
+
110
+ @dataclass
111
+ class AgentContext:
112
+ """
113
+ Built-in context for every BaseAgent instance
114
+ Automatically created when agent is instantiated
115
+ """
116
+ session_id: str
117
+ conversation_id: str
118
+ user_id: str
119
+ tenant_id: str
120
+ agent_id: str
121
+ created_at: datetime = field(default_factory=datetime.now)
122
+ metadata: Dict[str, Any] = field(default_factory=dict)
123
+
124
+ def to_execution_context(self) -> ExecutionContext:
125
+ """Convert to ExecutionContext for operations"""
126
+ return ExecutionContext(
127
+ session_id=self.session_id,
128
+ conversation_id=self.conversation_id,
129
+ user_id=self.user_id,
130
+ tenant_id=self.tenant_id,
131
+ metadata=self.metadata
132
+ )
133
+
134
+ def update_metadata(self, **kwargs):
135
+ """Update context metadata"""
136
+ self.metadata.update(kwargs)
137
+
138
+ def __str__(self):
139
+ return f"AgentContext(session={self.session_id}, user={self.user_id})"
140
+
141
+
142
+ @dataclass
143
+ class ProviderConfig:
144
+ """Configuration for LLM providers"""
145
+ name: str
146
+ model_name: str
147
+ priority: int
148
+ max_requests_per_minute: int = 60
149
+ max_requests_per_hour: int = 3600
150
+ cooldown_minutes: int = 5
151
+ request_count: int = 0
152
+ error_count: int = 0
153
+ last_request_time: Optional[datetime] = None
154
+ last_error_time: Optional[datetime] = None
155
+ is_available: bool = True
156
+
157
+ def __post_init__(self):
158
+ """Ensure no None values for numeric fields"""
159
+ if self.max_requests_per_minute is None:
160
+ self.max_requests_per_minute = 60
161
+ if self.max_requests_per_hour is None:
162
+ self.max_requests_per_hour = 3600
163
+ if self.request_count is None:
164
+ self.request_count = 0
165
+ if self.error_count is None:
166
+ self.error_count = 0
167
+ if self.priority is None:
168
+ self.priority = 999
169
+
170
+
171
+ class ProviderTracker:
172
+ """Tracks provider usage and availability"""
173
+
174
+ def __init__(self):
175
+ self.providers: Dict[str, ProviderConfig] = {}
176
+ self.current_provider: Optional[str] = None
177
+ self.last_rotation_time: Optional[datetime] = None
178
+ self.rotation_interval_minutes: int = 30
179
+
180
+ def record_request(self, provider_name: str):
181
+ """Record a request to a provider"""
182
+ if provider_name in self.providers:
183
+ provider = self.providers[provider_name]
184
+ provider.request_count += 1
185
+ provider.last_request_time = datetime.now()
186
+
187
+ def record_error(self, provider_name: str, error_message: str):
188
+ """Record an error for a provider"""
189
+ if provider_name in self.providers:
190
+ provider = self.providers[provider_name]
191
+ provider.error_count += 1
192
+ provider.last_error_time = datetime.now()
193
+
194
+ if provider.error_count >= 3:
195
+ provider.is_available = False
196
+
197
+ def is_provider_available(self, provider_name: str) -> bool:
198
+ """Check if a provider is available"""
199
+ if provider_name not in self.providers:
200
+ return False
201
+
202
+ provider = self.providers[provider_name]
203
+
204
+ if not provider.is_available:
205
+ if (provider.last_error_time and
206
+ datetime.now() - provider.last_error_time > timedelta(minutes=provider.cooldown_minutes)):
207
+ provider.is_available = True
208
+ provider.error_count = 0
209
+ else:
210
+ return False
211
+
212
+ now = datetime.now()
213
+ # FIXED: Check for None before arithmetic operations
214
+ if provider.last_request_time is not None:
215
+ time_since_last = (now - provider.last_request_time).total_seconds()
216
+ if time_since_last > 3600:
217
+ provider.request_count = 0
218
+
219
+ # FIXED: Ensure max_requests_per_hour is not None
220
+ max_requests = provider.max_requests_per_hour or 3600
221
+ if provider.request_count >= max_requests:
222
+ return False
223
+
224
+ return True
225
+
226
+ def get_best_available_provider(self) -> Optional[str]:
227
+ """Get the best available provider"""
228
+ available_providers = [
229
+ (name, config) for name, config in self.providers.items()
230
+ if self.is_provider_available(name)
231
+ ]
232
+
233
+ if not available_providers:
234
+ return None
235
+
236
+ def sort_key(provider_tuple):
237
+ name, config = provider_tuple
238
+ priority = config.priority or 999
239
+ error_count = config.error_count or 0
240
+ return (priority, error_count)
241
+
242
+ available_providers.sort(key=lambda x: (x[1].priority, x[1].error_count))
243
+ return available_providers[0][0]
244
+
245
+
246
+ class DockerCodeExecutor:
247
+ """Secure code execution using Docker containers"""
248
+
249
+ def __init__(self, config: Dict[str, Any]):
250
+ self.config = config
251
+ self.work_dir = config.get("work_dir", '/opt/ambivo/work_dir')
252
+ self.docker_images = config.get("docker_images", ["sgosain/amb-ubuntu-python-public-pod"])
253
+ self.timeout = config.get("timeout", 60)
254
+ self.default_image = self.docker_images[0] if self.docker_images else "sgosain/amb-ubuntu-python-public-pod"
255
+
256
+ if DOCKER_AVAILABLE:
257
+ try:
258
+ self.docker_client = docker.from_env()
259
+ self.docker_client.ping()
260
+ self.available = True
261
+ except Exception as e:
262
+ self.available = False
263
+ else:
264
+ self.available = False
265
+
266
+ def execute_code(self, code: str, language: str = "python", files: Dict[str, str] = None) -> Dict[str, Any]:
267
+ """Execute code in Docker container"""
268
+ if not self.available:
269
+ return {
270
+ 'success': False,
271
+ 'error': 'Docker not available',
272
+ 'language': language
273
+ }
274
+
275
+ try:
276
+ with tempfile.TemporaryDirectory() as temp_dir:
277
+ temp_path = Path(temp_dir)
278
+
279
+ if language == "python":
280
+ code_file = temp_path / "code.py"
281
+ code_file.write_text(code)
282
+ cmd = ["python", "/workspace/code.py"]
283
+ elif language == "bash":
284
+ code_file = temp_path / "script.sh"
285
+ code_file.write_text(code)
286
+ cmd = ["bash", "/workspace/script.sh"]
287
+ else:
288
+ raise ValueError(f"Unsupported language: {language}")
289
+
290
+ if files:
291
+ for filename, content in files.items():
292
+ file_path = temp_path / filename
293
+ file_path.write_text(content)
294
+
295
+ container_config = {
296
+ 'image': self.default_image,
297
+ 'command': cmd,
298
+ 'volumes': {str(temp_path): {'bind': '/workspace', 'mode': 'rw'}},
299
+ 'working_dir': '/workspace',
300
+ 'mem_limit': '512m',
301
+ 'network_disabled': True,
302
+ 'remove': True,
303
+ 'stdout': True,
304
+ 'stderr': True
305
+ }
306
+
307
+ start_time = time.time()
308
+ container = self.docker_client.containers.run(**container_config)
309
+ execution_time = time.time() - start_time
310
+
311
+ output = container.decode('utf-8') if isinstance(container, bytes) else str(container)
312
+
313
+ return {
314
+ 'success': True,
315
+ 'output': output,
316
+ 'execution_time': execution_time,
317
+ 'language': language
318
+ }
319
+
320
+ except docker.errors.ContainerError as e:
321
+ return {
322
+ 'success': False,
323
+ 'error': f"Container error: {e.stderr.decode('utf-8') if e.stderr else 'Unknown error'}",
324
+ 'exit_code': e.exit_status,
325
+ 'language': language
326
+ }
327
+ except Exception as e:
328
+ return {
329
+ 'success': False,
330
+ 'error': str(e),
331
+ 'language': language
332
+ }
333
+
334
+
335
+ class BaseAgent(ABC):
336
+ """
337
+ Enhanced BaseAgent with built-in auto-context session management and simplified chat interface
338
+ Every agent automatically gets a context with session_id, user_id, etc.
339
+ """
340
+
341
+ def __init__(self,
342
+ agent_id: str = None,
343
+ role: AgentRole = AgentRole.ASSISTANT,
344
+ user_id: str = None,
345
+ tenant_id: str = "default",
346
+ session_metadata: Dict[str, Any] = None,
347
+ memory_manager=None,
348
+ llm_service=None,
349
+ config: Dict[str, Any] = None,
350
+ name: str = None,
351
+ description: str = None,
352
+ auto_configure: bool = True,
353
+ session_id: str = None,
354
+ conversation_id: str = None,
355
+ **kwargs):
356
+
357
+ # Auto-generate agent_id if not provided
358
+ if agent_id is None:
359
+ agent_id = f"agent_{str(uuid.uuid4())[:8]}"
360
+
361
+ self.agent_id = agent_id
362
+ self.role = role
363
+ self.name = name or f"{role.value}_{agent_id[:8]}"
364
+ self.description = description or f"Agent with role: {role.value}"
365
+
366
+ # Load config if not provided and auto-configure is enabled
367
+ if config is None and auto_configure:
368
+ try:
369
+ from ..config.loader import load_config
370
+ config = load_config()
371
+ except Exception as e:
372
+ logging.warning(f"Could not load config for auto-configuration: {e}")
373
+ config = {}
374
+
375
+ self.config = config or {}
376
+
377
+ self.context = self._create_agent_context(user_id, tenant_id,
378
+ session_metadata,
379
+ session_id,
380
+ conversation_id)
381
+
382
+ # Auto-configure memory if not provided and auto-configure is enabled
383
+ if memory_manager is None and auto_configure:
384
+ try:
385
+ from ..core.memory import create_redis_memory_manager
386
+ self.memory = create_redis_memory_manager(
387
+ agent_id=agent_id,
388
+ redis_config=None # Will load from config automatically
389
+ )
390
+ logging.info(f"Auto-configured memory for agent {agent_id}")
391
+ except Exception as e:
392
+ logging.error(f"Failed to auto-configure memory for {agent_id}: {e}")
393
+ self.memory = None
394
+ else:
395
+ self.memory = memory_manager
396
+
397
+ # Auto-configure LLM service if not provided and auto-configure is enabled
398
+ if llm_service is None and auto_configure:
399
+ try:
400
+ from ..core.llm import create_multi_provider_llm_service
401
+ self.llm_service = create_multi_provider_llm_service()
402
+ logging.info(f"Auto-configured LLM service for agent {agent_id}")
403
+ except Exception as e:
404
+ logging.warning(f"Could not auto-configure LLM for {agent_id}: {e}")
405
+ self.llm_service = None
406
+ else:
407
+ self.llm_service = llm_service
408
+
409
+ self.tools = kwargs.get('tools', [])
410
+ self.active = True
411
+
412
+ # Initialize executor
413
+ self.executor = ThreadPoolExecutor(max_workers=4)
414
+
415
+ logging.info(f"🚀 BaseAgent created with auto-context:")
416
+ logging.info(f" 🤖 Agent: {self.agent_id}")
417
+ logging.info(f" 📋 Session: {self.context.session_id}")
418
+ logging.info(f" 👤 User: {self.context.user_id}")
419
+
420
+ def _create_agent_context(self,
421
+ user_id: str = None,
422
+ tenant_id: str = "default",
423
+ session_metadata: Dict[str, Any] = None,
424
+ session_id: str = None,
425
+ conversation_id: str = None
426
+ ) -> AgentContext:
427
+ """Create auto-context for this agent instance"""
428
+
429
+ # Auto-generate user_id if not provided
430
+ if user_id is None:
431
+ user_id = f"user_{str(uuid.uuid4())[:8]}"
432
+
433
+ if session_id and conversation_id:
434
+ final_session_id = session_id
435
+ final_conversation_id = conversation_id
436
+ else:
437
+ final_session_id = f"session_{str(uuid.uuid4())[:8]}"
438
+ final_conversation_id = f"conv_{str(uuid.uuid4())[:8]}"
439
+
440
+ return AgentContext(
441
+ session_id=final_session_id,
442
+ conversation_id=final_conversation_id,
443
+ user_id=user_id,
444
+ tenant_id=tenant_id,
445
+ agent_id=self.agent_id,
446
+ metadata=session_metadata or {}
447
+ )
448
+
449
+ @classmethod
450
+ def create(cls,
451
+ agent_id: str = None,
452
+ user_id: str = None,
453
+ tenant_id: str = "default",
454
+ session_metadata: Dict[str, Any] = None,
455
+ session_id: str = None,
456
+ conversation_id: str = None,
457
+ **kwargs) -> Tuple['BaseAgent', AgentContext]:
458
+ """
459
+ 🌟 DEFAULT: Create agent and return both agent and context
460
+ This is the RECOMMENDED way to create agents with auto-context
461
+
462
+ Usage:
463
+ agent, context = KnowledgeBaseAgent.create(user_id="john")
464
+ print(f"Session: {context.session_id}")
465
+ print(f"User: {context.user_id}")
466
+ """
467
+ if agent_id is None:
468
+ agent_id = f"{cls.__name__.lower()}_{str(uuid.uuid4())[:8]}"
469
+
470
+ agent = cls(
471
+ agent_id=agent_id,
472
+ user_id=user_id,
473
+ tenant_id=tenant_id,
474
+ session_metadata=session_metadata,
475
+ session_id=session_id,
476
+ conversation_id=conversation_id,
477
+ auto_configure=True,
478
+ **kwargs
479
+ )
480
+
481
+ return agent, agent.context
482
+
483
+ @classmethod
484
+ def create_simple(cls,
485
+ agent_id: str = None,
486
+ user_id: str = None,
487
+ tenant_id: str = "default",
488
+ session_metadata: Dict[str, Any] = None,
489
+ **kwargs) -> 'BaseAgent':
490
+ """
491
+ Create agent with auto-context (returns agent only)
492
+
493
+ ⚠️ LEGACY: Use create() instead for explicit context handling
494
+
495
+ Usage:
496
+ agent = KnowledgeBaseAgent.create_simple(user_id="john")
497
+ print(f"Session: {agent.context.session_id}") # Context still available
498
+ """
499
+ if agent_id is None:
500
+ agent_id = f"{cls.__name__.lower()}_{str(uuid.uuid4())[:8]}"
501
+
502
+ return cls(
503
+ agent_id=agent_id,
504
+ user_id=user_id,
505
+ tenant_id=tenant_id,
506
+ session_metadata=session_metadata,
507
+ auto_configure=True,
508
+ **kwargs
509
+ )
510
+
511
+ @classmethod
512
+ def create_advanced(cls,
513
+ agent_id: str,
514
+ memory_manager,
515
+ llm_service=None,
516
+ config: Dict[str, Any] = None,
517
+ user_id: str = None,
518
+ tenant_id: str = "default",
519
+ **kwargs):
520
+ """
521
+ Advanced factory method for explicit dependency injection
522
+
523
+ Usage:
524
+ memory = create_redis_memory_manager("custom_agent")
525
+ llm = create_multi_provider_llm_service()
526
+ agent = YouTubeDownloadAgent.create_advanced("my_id", memory, llm)
527
+ """
528
+ return cls(
529
+ agent_id=agent_id,
530
+ memory_manager=memory_manager,
531
+ llm_service=llm_service,
532
+ config=config,
533
+ user_id=user_id,
534
+ tenant_id=tenant_id,
535
+ auto_configure=False, # Disable auto-config when using advanced mode
536
+ **kwargs
537
+ )
538
+
539
+ # 🎯 NEW: SIMPLIFIED CHAT INTERFACE
540
+
541
+ async def chat(self, message: str, **kwargs) -> str:
542
+ """
543
+ 🌟 NEW: Simplified chat interface that converts string to full AgentMessage
544
+
545
+ This is the easiest way to interact with agents created via .create()
546
+ since they already have auto-context (session_id, user_id, etc.)
547
+
548
+ Args:
549
+ message: User message as string
550
+ **kwargs: Optional metadata to add to the message
551
+
552
+ Returns:
553
+ Agent response as string
554
+
555
+ Usage:
556
+ agent, context = YouTubeDownloadAgent.create(user_id="john")
557
+ response = await agent.chat("Download https://youtube.com/watch?v=abc123")
558
+ print(response)
559
+ """
560
+ try:
561
+ # Create AgentMessage from string using auto-context
562
+ user_message = AgentMessage(
563
+ id=str(uuid.uuid4()),
564
+ sender_id=self.context.user_id, # 🎯 Use auto-context user_id
565
+ recipient_id=self.agent_id,
566
+ content=message,
567
+ message_type=MessageType.USER_INPUT,
568
+ session_id=self.context.session_id, # 🎯 Use auto-context session_id
569
+ conversation_id=self.context.conversation_id, # 🎯 Use auto-context conversation_id
570
+ metadata={
571
+ 'chat_interface': True,
572
+ 'simplified_call': True,
573
+ **kwargs # Allow additional metadata
574
+ }
575
+ )
576
+
577
+ # Get execution context from auto-context
578
+ execution_context = self.context.to_execution_context()
579
+
580
+ # Add any additional metadata passed via kwargs
581
+ execution_context.metadata.update(kwargs)
582
+
583
+ # Call the full process_message method
584
+ agent_response = await self.process_message(user_message, execution_context)
585
+
586
+ # Return just the content string (simplified interface)
587
+ return agent_response.content
588
+
589
+ except Exception as e:
590
+ # Handle errors gracefully
591
+ error_msg = f"Chat error: {str(e)}"
592
+ logging.error(f"Agent {self.agent_id} chat error: {e}")
593
+ return error_msg
594
+
595
+ def chat_sync(self, message: str, **kwargs) -> str:
596
+ """
597
+ 🌟 NEW: Synchronous version of chat() for easier use in non-async contexts
598
+
599
+ Args:
600
+ message: User message as string
601
+ **kwargs: Optional metadata to add to the message
602
+
603
+ Returns:
604
+ Agent response as string
605
+
606
+ Usage:
607
+ agent, context = YouTubeDownloadAgent.create(user_id="john")
608
+ response = agent.chat_sync("Download https://youtube.com/watch?v=abc123")
609
+ print(response)
610
+ """
611
+ try:
612
+ # Get or create event loop
613
+ try:
614
+ loop = asyncio.get_event_loop()
615
+ if loop.is_running():
616
+ # If loop is already running, we need to use run_in_executor
617
+ import concurrent.futures
618
+ with concurrent.futures.ThreadPoolExecutor() as executor:
619
+ future = executor.submit(asyncio.run, self.chat(message, **kwargs))
620
+ return future.result()
621
+ else:
622
+ # Loop exists but not running
623
+ return loop.run_until_complete(self.chat(message, **kwargs))
624
+ except RuntimeError:
625
+ # No event loop exists, create one
626
+ return asyncio.run(self.chat(message, **kwargs))
627
+
628
+ except Exception as e:
629
+ error_msg = f"Sync chat error: {str(e)}"
630
+ logging.error(f"Agent {self.agent_id} sync chat error: {e}")
631
+ return error_msg
632
+
633
+ # 🔧 CONTEXT MANAGEMENT METHODS
634
+
635
+ def get_context(self) -> AgentContext:
636
+ """Get the agent's auto-generated context"""
637
+ return self.context
638
+
639
+ def get_execution_context(self) -> ExecutionContext:
640
+ """Get ExecutionContext for operations that need it"""
641
+ return self.context.to_execution_context()
642
+
643
+ def update_context_metadata(self, **kwargs):
644
+ """Update context metadata"""
645
+ self.context.update_metadata(**kwargs)
646
+
647
+ # 🧠 CONVERSATION HISTORY METHODS (Built into BaseAgent)
648
+
649
+ async def get_conversation_history(self,
650
+ limit: int = None,
651
+ include_metadata: bool = True) -> List[Dict[str, Any]]:
652
+ """
653
+ Get conversation history for this agent's session
654
+
655
+ Args:
656
+ limit: Maximum number of messages to return (None = all)
657
+ include_metadata: Whether to include message metadata
658
+
659
+ Returns:
660
+ List of conversation messages with context
661
+ """
662
+ try:
663
+ if not self.memory:
664
+ logging.warning(f"No memory available for agent {self.agent_id}")
665
+ return []
666
+
667
+ # Get history using session_id from auto-context
668
+ history = self.memory.get_recent_messages(
669
+ limit=limit or 10,
670
+ conversation_id=self.context.conversation_id
671
+ )
672
+
673
+ # Add context information to each message
674
+ enriched_history = []
675
+ for msg in history:
676
+ if include_metadata:
677
+ msg_with_context = {
678
+ **msg,
679
+ 'session_id': self.context.session_id,
680
+ 'user_id': self.context.user_id,
681
+ 'agent_id': self.agent_id,
682
+ 'conversation_id': self.context.conversation_id
683
+ }
684
+ else:
685
+ msg_with_context = msg
686
+
687
+ enriched_history.append(msg_with_context)
688
+
689
+ return enriched_history
690
+
691
+ except Exception as e:
692
+ logging.error(f"Failed to get conversation history for {self.agent_id}: {e}")
693
+ return []
694
+
695
+ async def add_to_conversation_history(self,
696
+ message: str,
697
+ message_type: str = "user",
698
+ metadata: Dict[str, Any] = None) -> bool:
699
+ """
700
+ Add a message to conversation history
701
+
702
+ Args:
703
+ message: The message content
704
+ message_type: Type of message ("user", "agent", "system")
705
+ metadata: Additional metadata for the message
706
+
707
+ Returns:
708
+ True if successfully added, False otherwise
709
+ """
710
+ try:
711
+ if not self.memory:
712
+ logging.warning(f"No memory available for agent {self.agent_id}")
713
+ return False
714
+
715
+ # Create AgentMessage for storage
716
+ agent_message = AgentMessage(
717
+ id=str(uuid.uuid4()),
718
+ sender_id=self.agent_id if message_type == "agent" else f"{message_type}_sender",
719
+ recipient_id=None,
720
+ content=message,
721
+ message_type=MessageType.AGENT_RESPONSE if message_type == "agent" else MessageType.USER_INPUT,
722
+ session_id=self.context.session_id,
723
+ conversation_id=self.context.conversation_id,
724
+ metadata={
725
+ 'type': message_type,
726
+ 'user_id': self.context.user_id,
727
+ 'agent_id': self.agent_id,
728
+ **(metadata or {})
729
+ }
730
+ )
731
+
732
+ # Store in memory
733
+ self.memory.store_message(agent_message)
734
+ return True
735
+
736
+ except Exception as e:
737
+ logging.error(f"Failed to add to conversation history for {self.agent_id}: {e}")
738
+ return False
739
+
740
+ async def clear_conversation_history(self) -> bool:
741
+ """
742
+ Clear conversation history for this agent's session
743
+
744
+ Returns:
745
+ True if successfully cleared, False otherwise
746
+ """
747
+ try:
748
+ if not self.memory:
749
+ logging.warning(f"No memory available for agent {self.agent_id}")
750
+ return False
751
+
752
+ self.memory.clear_memory(self.context.conversation_id)
753
+ logging.info(f"Cleared conversation history for session {self.context.session_id}")
754
+ return True
755
+
756
+ except Exception as e:
757
+ logging.error(f"Failed to clear conversation history for {self.agent_id}: {e}")
758
+ return False
759
+
760
+ async def get_conversation_summary(self) -> Dict[str, Any]:
761
+ """
762
+ Get a summary of the current conversation
763
+
764
+ Returns:
765
+ Dictionary with conversation statistics and summary
766
+ """
767
+ try:
768
+ history = await self.get_conversation_history(include_metadata=True)
769
+
770
+ if not history:
771
+ return {
772
+ 'total_messages': 0,
773
+ 'user_messages': 0,
774
+ 'agent_messages': 0,
775
+ 'session_duration': '0 minutes',
776
+ 'first_message': None,
777
+ 'last_message': None,
778
+ 'session_id': self.context.session_id
779
+ }
780
+
781
+ # Analyze conversation
782
+ total_messages = len(history)
783
+ user_messages = len([msg for msg in history if msg.get('message_type') == 'user_input'])
784
+ agent_messages = len([msg for msg in history if msg.get('message_type') == 'agent_response'])
785
+
786
+ # Calculate session duration
787
+ first_msg_time = self.context.created_at
788
+ last_msg_time = datetime.now()
789
+ duration = last_msg_time - first_msg_time
790
+ duration_minutes = int(duration.total_seconds() / 60)
791
+
792
+ return {
793
+ 'total_messages': total_messages,
794
+ 'user_messages': user_messages,
795
+ 'agent_messages': agent_messages,
796
+ 'session_duration': f"{duration_minutes} minutes",
797
+ 'first_message': history[0].get('content', '')[:100] + "..." if len(
798
+ history[0].get('content', '')) > 100 else history[0].get('content', '') if history else None,
799
+ 'last_message': history[-1].get('content', '')[:100] + "..." if len(
800
+ history[-1].get('content', '')) > 100 else history[-1].get('content', '') if history else None,
801
+ 'session_id': self.context.session_id,
802
+ 'conversation_id': self.context.conversation_id,
803
+ 'user_id': self.context.user_id
804
+ }
805
+
806
+ except Exception as e:
807
+ logging.error(f"Failed to get conversation summary for {self.agent_id}: {e}")
808
+ return {
809
+ 'error': str(e),
810
+ 'session_id': self.context.session_id
811
+ }
812
+
813
+ async def _with_auto_context(self, operation_name: str, **kwargs) -> Dict[str, Any]:
814
+ """
815
+ Internal method that automatically applies context to operations
816
+ All agent operations should use this to ensure context is applied
817
+ """
818
+ execution_context = self.get_execution_context()
819
+
820
+ # Add context info to operation metadata
821
+ operation_metadata = {
822
+ 'session_id': self.context.session_id,
823
+ 'user_id': self.context.user_id,
824
+ 'tenant_id': self.context.tenant_id,
825
+ 'operation': operation_name,
826
+ 'timestamp': datetime.now().isoformat(),
827
+ **kwargs
828
+ }
829
+
830
+ # Update context metadata
831
+ self.context.update_metadata(**operation_metadata)
832
+
833
+ return {
834
+ 'execution_context': execution_context,
835
+ 'operation_metadata': operation_metadata
836
+ }
837
+
838
+ # 🧹 SESSION CLEANUP
839
+
840
+ async def cleanup_session(self) -> bool:
841
+ """Cleanup the agent's session and resources"""
842
+ try:
843
+ session_id = self.context.session_id
844
+
845
+ # Clear memory for this session
846
+ if hasattr(self, 'memory') and self.memory:
847
+ try:
848
+ # Commented out temporarily as noted in original
849
+ # self.memory.clear_memory(self.context.conversation_id)
850
+ logging.info(f"🧹 Cleared memory for session {session_id}")
851
+ except Exception as e:
852
+ logging.warning(f"⚠️ Could not clear memory: {e}")
853
+
854
+ # Shutdown executor
855
+ if hasattr(self, 'executor') and self.executor:
856
+ try:
857
+ self.executor.shutdown(wait=True)
858
+ logging.info(f"🛑 Shutdown executor for session {session_id}")
859
+ except Exception as e:
860
+ logging.warning(f"⚠️ Could not shutdown executor: {e}")
861
+
862
+ logging.info(f"✅ Session {session_id} cleaned up successfully")
863
+ return True
864
+
865
+ except Exception as e:
866
+ logging.error(f"❌ Error cleaning up session: {e}")
867
+ return False
868
+
869
+ # 🛠️ TOOL MANAGEMENT
870
+
871
+ def add_tool(self, tool: AgentTool):
872
+ """Add a tool to the agent"""
873
+ self.tools.append(tool)
874
+
875
+ def get_tool(self, tool_name: str) -> Optional[AgentTool]:
876
+ """Get a tool by name"""
877
+ return next((tool for tool in self.tools if tool.name == tool_name), None)
878
+
879
+ async def execute_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
880
+ """Execute a tool with auto-context"""
881
+ tool = self.get_tool(tool_name)
882
+ if not tool:
883
+ raise ValueError(f"Tool {tool_name} not found")
884
+
885
+ # Apply auto-context to tool execution
886
+ context_data = await self._with_auto_context("tool_execution",
887
+ tool_name=tool_name,
888
+ parameters=parameters)
889
+
890
+ try:
891
+ if asyncio.iscoroutinefunction(tool.function):
892
+ result = await tool.function(**parameters)
893
+ else:
894
+ result = await asyncio.get_event_loop().run_in_executor(
895
+ self.executor, tool.function, **parameters
896
+ )
897
+
898
+ return {
899
+ 'success': True,
900
+ 'result': result,
901
+ 'session_id': self.context.session_id,
902
+ 'context': context_data
903
+ }
904
+ except Exception as e:
905
+ return {
906
+ 'success': False,
907
+ 'error': str(e),
908
+ 'session_id': self.context.session_id
909
+ }
910
+
911
+ def create_response(self,
912
+ content: str,
913
+ recipient_id: str,
914
+ message_type: MessageType = MessageType.AGENT_RESPONSE,
915
+ metadata: Dict[str, Any] = None,
916
+ session_id: str = None,
917
+ conversation_id: str = None) -> AgentMessage:
918
+ """
919
+ Create a response message with auto-context
920
+ Uses agent's context if session_id/conversation_id not provided
921
+ """
922
+ return AgentMessage(
923
+ id=str(uuid.uuid4()),
924
+ sender_id=self.agent_id,
925
+ recipient_id=recipient_id,
926
+ content=content,
927
+ message_type=message_type,
928
+ metadata=metadata or {},
929
+ session_id=session_id or self.context.session_id, # 🎯 Auto-context!
930
+ conversation_id=conversation_id or self.context.conversation_id # 🎯 Auto-context!
931
+ )
932
+
933
+ # 📨 ABSTRACT METHOD (must be implemented by subclasses)
934
+
935
+ @abstractmethod
936
+ async def process_message(self, message: AgentMessage, context: ExecutionContext = None) -> AgentMessage:
937
+ """
938
+ Process incoming message and return response
939
+ Uses agent's auto-context if context not provided
940
+ """
941
+ if context is None:
942
+ context = self.get_execution_context()
943
+
944
+ # Subclasses must implement this
945
+ pass
946
+
947
+ def register_agent(self, agent: 'BaseAgent'):
948
+ """Default implementation - only ProxyAgent should override this"""
949
+ return False
950
+
951
+
952
+ # 🎯 CONTEXT MANAGER FOR AUTO-CONTEXT AGENTS
953
+
954
+ class AgentSession:
955
+ """
956
+ Context manager for BaseAgent instances with automatic cleanup
957
+
958
+ Usage:
959
+ async with AgentSession(KnowledgeBaseAgent, user_id="john") as agent:
960
+ result = await agent.chat("What is machine learning?")
961
+ print(f"Session: {agent.context.session_id}")
962
+ # Agent automatically cleaned up
963
+ """
964
+
965
+ def __init__(self,
966
+ agent_class,
967
+ user_id: str = None,
968
+ tenant_id: str = "default",
969
+ session_metadata: Dict[str, Any] = None,
970
+ **agent_kwargs):
971
+ self.agent_class = agent_class
972
+ self.user_id = user_id
973
+ self.tenant_id = tenant_id
974
+ self.session_metadata = session_metadata
975
+ self.agent_kwargs = agent_kwargs
976
+ self.agent = None
977
+
978
+ async def __aenter__(self):
979
+ """Create agent when entering context"""
980
+ self.agent = self.agent_class.create_simple(
981
+ user_id=self.user_id,
982
+ tenant_id=self.tenant_id,
983
+ session_metadata=self.session_metadata,
984
+ **self.agent_kwargs
985
+ )
986
+ return self.agent
987
+
988
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
989
+ """Cleanup agent when exiting context"""
990
+ if self.agent:
991
+ await self.agent.cleanup_session()
992
+
993
+
994
+ # 🚀 CONVENIENCE FUNCTIONS FOR QUICK AGENT USAGE
995
+
996
+ async def quick_chat(agent_class, message: str, user_id: str = None, **kwargs) -> str:
997
+ """
998
+ 🌟 ULTRA-SIMPLIFIED: One-liner agent chat
999
+
1000
+ Usage:
1001
+ response = await quick_chat(YouTubeDownloadAgent, "Download https://youtube.com/watch?v=abc")
1002
+ print(response)
1003
+ """
1004
+ try:
1005
+ agent = agent_class.create_simple(user_id=user_id, **kwargs)
1006
+ response = await agent.chat(message)
1007
+ await agent.cleanup_session()
1008
+ return response
1009
+ except Exception as e:
1010
+ return f"Quick chat error: {str(e)}"
1011
+
1012
+
1013
+ def quick_chat_sync(agent_class, message: str, user_id: str = None, **kwargs) -> str:
1014
+ """
1015
+ 🌟 ULTRA-SIMPLIFIED: One-liner synchronous agent chat
1016
+
1017
+ Usage:
1018
+ response = quick_chat_sync(YouTubeDownloadAgent, "Download https://youtube.com/watch?v=abc")
1019
+ print(response)
1020
+ """
1021
+ try:
1022
+ return asyncio.run(quick_chat(agent_class, message, user_id, **kwargs))
1023
+ except Exception as e:
1024
+ return f"Quick sync chat error: {str(e)}"