fix: work with async agents

Author: k11kirky

agentserve/__init__.py

@@ -1,2 +1,5 @@
 # agentserve/__init__.py
 from .agent_server import AgentServer as app
+from .logging_config import setup_logger
+
+logger = setup_logger()

agentserve/agent_registry.py

@@ -1,30 +1,45 @@
 # agentserve/agent_registry.py
 from typing import Callable, Optional, Type
 from pydantic import BaseModel
+from .logging_config import setup_logger
+import asyncio
 
 class AgentRegistry:
     def __init__(self):
         self.agent_function = None
         self.input_schema: Optional[Type[BaseModel]] = None
-
+        self.logger = setup_logger("agentserve.agent_registry")
+        
     def register_agent(self, func: Optional[Callable] = None, *, input_schema: Optional[Type[BaseModel]] = None):
         if func is None:
-            # Decorator is called with arguments
             def wrapper(func: Callable):
                 return self.register_agent(func, input_schema=input_schema)
             return wrapper
 
         self.input_schema = input_schema
+        is_async = asyncio.iscoroutinefunction(func)
+        self.logger.info(f"Registering {'async' if is_async else 'sync'} function")
 
-        def validated_func(task_data):
+        async def async_validated_func(task_data):
+            if self.input_schema is not None:
+                validated_data = self.input_schema(**task_data)
+                return await func(validated_data)
+            return await func(task_data)
+
+        def sync_validated_func(task_data):
             if self.input_schema is not None:
                 validated_data = self.input_schema(**task_data)
                 return func(validated_data)
-            else:
-                return func(task_data)
+            return func(task_data)
+
+        if is_async:
+            self.agent_function = async_validated_func
+            setattr(self.agent_function, '_is_async', True)
+        else:
+            self.agent_function = sync_validated_func
+            setattr(self.agent_function, '_is_async', False)
 
-        self.agent_function = validated_func
-        return validated_func
+        return self.agent_function
 
     def get_agent(self):
         if self.agent_function is None:

agentserve/agent_server.py

@@ -6,35 +6,50 @@
 from .agent_registry import AgentRegistry
 from typing import Dict, Any, Optional
 from .config import Config
+from .logging_config import setup_logger
 import uuid
 
 class AgentServer:
     def __init__(self, config: Optional[Config] = None):
-        self.app = FastAPI()
+        self.logger = setup_logger("agentserve.server")
+        self.app = FastAPI(debug=True)
         self.agent_registry = AgentRegistry()
         self.config = config or Config()
         self.task_queue = self._initialize_task_queue()
         self.agent = self.agent_registry.register_agent
         self._setup_routes()
+        self.logger.info("AgentServer initialized")
 
     def _initialize_task_queue(self):
         task_queue_type = self.config.get('task_queue', 'local').lower()
-        if task_queue_type == 'celery':
-            from .celery_task_queue import CeleryTaskQueue
-            return CeleryTaskQueue(self.config)
-        elif task_queue_type == 'redis':
-            from .redis_task_queue import RedisTaskQueue
-            return RedisTaskQueue(self.config)
-        else:
-            from .queues.local_task_queue import LocalTaskQueue
-            return LocalTaskQueue()
+        self.logger.info(f"Initializing {task_queue_type} task queue")
+        
+        try:
+            if task_queue_type == 'celery':
+                from .queues.celery_task_queue import CeleryTaskQueue
+                return CeleryTaskQueue(self.config)
+            elif task_queue_type == 'redis':
+                from .queues.redis_task_queue import RedisTaskQueue
+                return RedisTaskQueue(self.config)
+            else:
+                from .queues.local_task_queue import LocalTaskQueue
+                return LocalTaskQueue()
+        except Exception as e:
+            self.logger.error(f"Failed to initialize task queue: {str(e)}")
+            raise
 
     def _setup_routes(self):
         @self.app.post("/task/sync")
         async def sync_task(task_data: Dict[str, Any]):
+            self.logger.debug(f"sync_task called with data: {task_data}")
             try:
                 agent_function = self.agent_registry.get_agent()
-                result = agent_function(task_data)
+                if getattr(agent_function, '_is_async', False):
+                    self.logger.info("Function is async, running in event loop")
+                    result = await agent_function(task_data)
+                else:
+                    self.logger.info("Function is sync, running directly")
+                    result = agent_function(task_data)
                 return {"result": result}
             except ValidationError as ve:
                 if hasattr(ve, 'errors'):

agentserve/logging_config.py

@@ -0,0 +1,19 @@
+import logging
+import sys
+from typing import Optional
+
+def setup_logger(name: str = "agentserve", level: Optional[str] = None) -> logging.Logger:
+    logger = logging.getLogger(name)
+    
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        formatter = logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+    
+    log_level = getattr(logging, (level or "DEBUG").upper())
+    logger.setLevel(log_level)
+    
+    return logger

agentserve/queues/calery_task_queue.py

@@ -1,30 +1,44 @@
 # agentserve/celery_task_queue.py
 
+import asyncio
 from typing import Any, Dict
 from .task_queue import TaskQueue
 from ..config import Config
+from ..logging_config import setup_logger
+
 class CeleryTaskQueue(TaskQueue):
     def __init__(self, config: Config):
         try:
             from celery import Celery
         except ImportError:
             raise ImportError("CeleryTaskQueue requires the 'celery' package. Please install it.")
 
+        self.logger = setup_logger("agentserve.queue.celery")
         broker_url = config.get('celery', {}).get('broker_url', 'pyamqp://guest@localhost//')
         self.celery_app = Celery('agent_server', broker=broker_url)
+        self.loop = asyncio.new_event_loop()
         self._register_tasks()
+        self.logger.info("CeleryTaskQueue initialized")
 
     def _register_tasks(self):
         @self.celery_app.task(name='agent_task')
-        def agent_task(task_data):
-            from .agent_registry import AgentRegistry
+        def agent_task(task_data, is_async=False):
+            from ..agent_registry import AgentRegistry
             agent_registry = AgentRegistry()
             agent_function = agent_registry.get_agent()
+            
+            if is_async:
+                asyncio.set_event_loop(self.loop)
+                return self.loop.run_until_complete(agent_function(task_data))
             return agent_function(task_data)
 
     def enqueue(self, agent_function, task_data: Dict[str, Any], task_id: str):
-        # Since the agent task is registered with Celery, we just send the task name
-        self.celery_app.send_task('agent_task', args=[task_data], task_id=task_id)
+        self.logger.debug(f"Enqueueing task {task_id}")
+        is_async = getattr(agent_function, '_is_async', False)
+        self.celery_app.send_task('agent_task', 
+                                args=[task_data], 
+                                kwargs={'is_async': is_async}, 
+                                task_id=task_id)
 
     def get_status(self, task_id: str) -> str:
         result = self.celery_app.AsyncResult(task_id)

agentserve/queues/local_task_queue.py

@@ -4,23 +4,35 @@
 from typing import Any, Dict
 from .task_queue import TaskQueue
 import threading
+from ..logging_config import setup_logger
 
 class LocalTaskQueue(TaskQueue):
     def __init__(self):
+        self.logger = setup_logger("agentserve.queue.local")
         self.results = {}
         self.statuses = {}
+        self.loop = asyncio.new_event_loop()
+        self.logger.info("LocalTaskQueue initialized")
 
     def enqueue(self, agent_function, task_data: Dict[str, Any], task_id: str):
+        self.logger.debug(f"Enqueueing task {task_id}")
         self.statuses[task_id] = 'queued'
         threading.Thread(target=self._run_task, args=(agent_function, task_data, task_id)).start()
 
     def _run_task(self, agent_function, task_data: Dict[str, Any], task_id: str):
+        self.logger.debug(f"Starting task {task_id}")
         self.statuses[task_id] = 'in_progress'
         try:
-            result = agent_function(task_data)
+            if getattr(agent_function, '_is_async', False):
+                asyncio.set_event_loop(self.loop)
+                result = self.loop.run_until_complete(agent_function(task_data))
+            else:
+                result = agent_function(task_data)
             self.results[task_id] = result
             self.statuses[task_id] = 'completed'
+            self.logger.info(f"Task {task_id} completed successfully")
         except Exception as e:
+            self.logger.error(f"Task {task_id} failed: {str(e)}")
             self.results[task_id] = e
             self.statuses[task_id] = 'failed'
 

agentserve/queues/redis_task_queue.py

@@ -1,7 +1,9 @@
 # agentserve/redis_task_queue.py
 
+import asyncio
 from typing import Any, Dict
 from .task_queue import TaskQueue
+from ..logging_config import setup_logger
 
 class RedisTaskQueue(TaskQueue):
     def __init__(self, config: Config):
@@ -11,14 +13,28 @@ def __init__(self, config: Config):
         except ImportError:
             raise ImportError("RedisTaskQueue requires 'redis' and 'rq' packages. Please install them.")
 
+        self.logger = setup_logger("agentserve.queue.redis")
         redis_config = config.get('redis', {})
         redis_host = redis_config.get('host', 'localhost')
         redis_port = redis_config.get('port', 6379)
         self.redis_conn = Redis(host=redis_host, port=redis_port)
         self.task_queue = Queue(connection=self.redis_conn)
+        self.loop = asyncio.new_event_loop()
+        self.logger.info("RedisTaskQueue initialized")
 
     def enqueue(self, agent_function, task_data: Dict[str, Any], task_id: str):
-        self.task_queue.enqueue_call(func=agent_function, args=(task_data,), job_id=task_id)
+        self.logger.debug(f"Enqueueing task {task_id}")
+        if getattr(agent_function, '_is_async', False):
+            wrapped_func = self._wrap_async_function(agent_function)
+            self.task_queue.enqueue_call(func=wrapped_func, args=(task_data,), job_id=task_id)
+        else:
+            self.task_queue.enqueue_call(func=agent_function, args=(task_data,), job_id=task_id)
+    
+    def _wrap_async_function(self, func):
+        def wrapper(task_data):
+            asyncio.set_event_loop(self.loop)
+            return self.loop.run_until_complete(func(task_data))
+        return wrapper
 
     def get_status(self, task_id: str) -> str:
         job = self.task_queue.fetch_job(task_id)

async_example.py

@@ -0,0 +1,18 @@
+import agentserve
+from pydantic import BaseModel
+import asyncio
+
+# Configure logging level
+agentserve.setup_logger(level="DEBUG")  # or "INFO", "WARNING", "ERROR"
+
+app = agentserve.app()
+
+class MyInputSchema(BaseModel):
+    prompt: str
+
+@app.agent(input_schema=MyInputSchema)
+async def my_agent(task_data):
+    await asyncio.sleep(1)
+    return task_data
+
+app.run()

example.py

@@ -1,7 +1,12 @@
 import agentserve
 from pydantic import BaseModel
+
+# Configure logging level
+agentserve.setup_logger(level="DEBUG")  # or "INFO", "WARNING", "ERROR"
+
 app = agentserve.app()
 
+
 class MyInputSchema(BaseModel):
     prompt: str