[WIP] add agentic RAG notebook (#21)

llama_agents/__init__.py

@@ -4,12 +4,24 @@
 from llama_agents.orchestrators import (
     AgentOrchestrator,
     PipelineOrchestrator,
-    ServiceComponent,
-    ServiceTool,
 )
-from llama_agents.tools import MetaServiceTool
+from llama_agents.tools import MetaServiceTool, ServiceComponent, ServiceTool
 from llama_agents.services import AgentService, ToolService, HumanService
 
+# configure logger
+import logging
+
+root_logger = logging.getLogger("llama_agents")
+
+formatter = logging.Formatter("%(levelname)s:%(name)s - %(message)s")
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(formatter)
+root_logger.addHandler(console_handler)
+
+root_logger.setLevel(logging.WARNING)
+root_logger.propagate = False
+
+
 __all__ = [
     # services
     "AgentService",

llama_agents/control_plane/server.py

@@ -1,6 +1,7 @@
 import uuid
 import uvicorn
 from fastapi import FastAPI
+from logging import getLogger
 from typing import Dict, List, Optional
 
 from llama_index.core import StorageContext, VectorStoreIndex
@@ -16,19 +17,15 @@
 from llama_agents.message_queues.base import BaseMessageQueue, PublishCallback
 from llama_agents.messages.base import QueueMessage
 from llama_agents.orchestrators.base import BaseOrchestrator
-from llama_agents.orchestrators.service_tool import ServiceTool
+from llama_agents.tools import ServiceTool
 from llama_agents.types import (
     ActionTypes,
     ServiceDefinition,
     TaskDefinition,
     TaskResult,
 )
 
-import logging
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-logging.basicConfig(level=logging.INFO)
+logger = getLogger(__name__)
 
 
 class ControlPlaneServer(BaseControlPlane):

llama_agents/launchers/local.py

@@ -83,6 +83,9 @@ def signal_handler(sig: Any, frame: Any) -> None:
         return signal_handler
 
     async def alaunch_single(self, initial_task: str) -> str:
+        # clear any result
+        self.result = None
+
         # register human consumer
         human_consumer = HumanMessageConsumer(
             message_handler={
@@ -98,7 +101,9 @@ async def alaunch_single(self, initial_task: str) -> str:
         # start services
         bg_tasks: List[asyncio.Task] = []
         for service in self.services:
-            bg_tasks.append(asyncio.create_task(service.launch_local()))
+            if hasattr(service, "raise_exceptions"):
+                service.raise_exceptions = True  # ensure exceptions are raised
+            bg_tasks.append(await service.launch_local())
 
         # publish initial task
         await self.publish(
@@ -109,7 +114,7 @@ async def alaunch_single(self, initial_task: str) -> str:
             ),
         )
         # runs until the message queue is stopped by the human consumer
-        mq_task = asyncio.create_task(self.message_queue.launch_local())
+        mq_task = await self.message_queue.launch_local()
         shutdown_handler = self.get_shutdown_handler([mq_task] + bg_tasks)
         loop = asyncio.get_event_loop()
         while loop.is_running():

llama_agents/message_queues/__init__.py

@@ -1,4 +1,7 @@
 from llama_agents.message_queues.base import BaseMessageQueue
-from llama_agents.message_queues.simple import SimpleMessageQueue
+from llama_agents.message_queues.simple import (
+    SimpleMessageQueue,
+    SimpleRemoteClientMessageQueue,
+)
 
-__all__ = ["BaseMessageQueue", "SimpleMessageQueue"]
+__all__ = ["BaseMessageQueue", "SimpleMessageQueue", "SimpleRemoteClientMessageQueue"]

llama_agents/message_queues/base.py

@@ -1,7 +1,9 @@
 """Message queue module."""
-
+import asyncio
 import inspect
+
 from abc import ABC, abstractmethod
+from logging import getLogger
 from pydantic import BaseModel
 from typing import Any, List, Optional, Protocol, TYPE_CHECKING
 
@@ -10,11 +12,7 @@
 if TYPE_CHECKING:
     from llama_agents.message_consumers.base import BaseMessageQueueConsumer
 
-import logging
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-logging.basicConfig(level=logging.INFO)
+logger = getLogger(__name__)
 
 
 class MessageProcessor(Protocol):
@@ -49,10 +47,14 @@ async def publish(
         self,
         message: QueueMessage,
         callback: Optional[PublishCallback] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> Any:
         """Send message to a consumer."""
-        logger.info("Publishing message: " + str(message))
+        logger.info(
+            f"Publishing message to '{message.type}' with action '{message.action}'"
+        )
+        logger.debug(f"Message: {message.model_dump()}")
+
         message.stats.publish_time = message.stats.timestamp_str()
         await self._publish(message)
 
@@ -88,7 +90,7 @@ async def processing_loop(self) -> None:
         ...
 
     @abstractmethod
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         """Launch the service in-process."""
         ...
 

llama_agents/message_queues/simple.py

@@ -3,12 +3,12 @@
 import asyncio
 import httpx
 import random
-import logging
 import uvicorn
 
 from collections import deque
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
+from logging import getLogger
 from pydantic import Field, PrivateAttr
 from typing import Any, AsyncGenerator, Dict, List, Optional
 from urllib.parse import urljoin
@@ -22,9 +22,7 @@
 )
 from llama_agents.types import PydanticValidatedUrl
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-logging.basicConfig(level=logging.DEBUG)
+logger = getLogger(__name__)
 
 
 class SimpleRemoteClientMessageQueue(BaseMessageQueue):
@@ -101,7 +99,7 @@ async def processing_loop(self) -> None:
             "`procesing_loop()` is not implemented for this class."
         )
 
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         raise NotImplementedError("`launch_local()` is not implemented for this class.")
 
     async def launch_server(self) -> None:
@@ -283,9 +281,9 @@ async def lifespan(self, app: FastAPI) -> AsyncGenerator[None, None]:
         yield
         self.running = False
 
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         logger.info("Launching message queue locally")
-        asyncio.create_task(self.processing_loop())
+        return asyncio.create_task(self.processing_loop())
 
     async def launch_server(self) -> None:
         logger.info(f"Launching message queue server at {self.host}:{self.port}")

llama_agents/orchestrators/__init__.py

@@ -1,13 +1,9 @@
 from llama_agents.orchestrators.agent import AgentOrchestrator
 from llama_agents.orchestrators.base import BaseOrchestrator
 from llama_agents.orchestrators.pipeline import PipelineOrchestrator
-from llama_agents.orchestrators.service_component import ServiceComponent
-from llama_agents.orchestrators.service_tool import ServiceTool
 
 __all__ = [
     "BaseOrchestrator",
     "PipelineOrchestrator",
-    "ServiceComponent",
-    "ServiceTool",
     "AgentOrchestrator",
 ]

llama_agents/orchestrators/agent.py

@@ -6,7 +6,7 @@
 
 from llama_agents.messages.base import QueueMessage
 from llama_agents.orchestrators.base import BaseOrchestrator
-from llama_agents.orchestrators.service_tool import ServiceTool
+from llama_agents.tools.service_tool import ServiceTool
 from llama_agents.types import ActionTypes, ChatMessage, TaskDefinition, TaskResult
 
 HISTORY_KEY = "chat_history"

llama_agents/orchestrators/pipeline.py

@@ -7,7 +7,7 @@
 
 from llama_agents.messages.base import QueueMessage
 from llama_agents.orchestrators.base import BaseOrchestrator
-from llama_agents.orchestrators.service_component import ServiceComponent
+from llama_agents.tools.service_component import ServiceComponent
 from llama_agents.types import ActionTypes, TaskDefinition, TaskResult
 
 RUN_STATE_KEY = "run_state"

llama_agents/services/agent.py

@@ -3,6 +3,7 @@
 import uvicorn
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
+from logging import getLogger
 from pydantic import PrivateAttr
 from typing import AsyncGenerator, Dict, List, Literal, Optional
 
@@ -25,11 +26,7 @@
     CONTROL_PLANE_NAME,
 )
 
-import logging
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-logging.basicConfig(level=logging.DEBUG)
+logger = getLogger(__name__)
 
 
 class AgentService(BaseService):
@@ -41,6 +38,7 @@ class AgentService(BaseService):
     step_interval: float = 0.1
     host: Optional[str] = None
     port: Optional[int] = None
+    raise_exceptions: bool = False
 
     _message_queue: BaseMessageQueue = PrivateAttr()
     _app: FastAPI = PrivateAttr()
@@ -59,6 +57,7 @@ def __init__(
         step_interval: float = 0.1,
         host: Optional[str] = None,
         port: Optional[int] = None,
+        raise_exceptions: bool = False,
     ) -> None:
         super().__init__(
             agent=agent,
@@ -69,6 +68,7 @@ def __init__(
             prompt=prompt,
             host=host,
             port=port,
+            raise_exceptions=raise_exceptions,
         )
 
         self._message_queue = message_queue
@@ -173,6 +173,13 @@ async def processing_loop(self) -> None:
                         )
             except Exception as e:
                 logger.error(f"Error in {self.service_name} processing_loop: {e}")
+                if self.raise_exceptions:
+                    # Kill everything
+                    # TODO: is there a better way to do this?
+                    import signal
+
+                    signal.raise_signal(signal.SIGINT)
+
                 continue
 
             await asyncio.sleep(self.step_interval)
@@ -197,9 +204,9 @@ def as_consumer(self, remote: bool = False) -> BaseMessageQueueConsumer:
             handler=self.process_message,
         )
 
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         logger.info(f"{self.service_name} launch_local")
-        asyncio.create_task(self.processing_loop())
+        return asyncio.create_task(self.processing_loop())
 
     # ---- Server based methods ----
 

llama_agents/services/base.py

@@ -1,3 +1,4 @@
+import asyncio
 import httpx
 from abc import ABC, abstractmethod
 from pydantic import BaseModel
@@ -46,7 +47,7 @@ async def publish(self, message: QueueMessage, **kwargs: Any) -> None:
         )
 
     @abstractmethod
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         """Launch the service in-process."""
         ...
 

llama_agents/services/human.py

@@ -1,9 +1,9 @@
 import asyncio
-import logging
 import uuid
 import uvicorn
 from asyncio import Lock
 from fastapi import FastAPI
+from logging import getLogger
 from pydantic import PrivateAttr
 from typing import Dict, List, Optional
 
@@ -27,9 +27,7 @@
 )
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-logging.basicConfig(level=logging.DEBUG)
+logger = getLogger(__name__)
 
 
 HELP_REQUEST_TEMPLATE_STR = (
@@ -200,9 +198,9 @@ def as_consumer(self, remote: bool = False) -> BaseMessageQueueConsumer:
             handler=self.process_message,
         )
 
-    async def launch_local(self) -> None:
+    async def launch_local(self) -> asyncio.Task:
         logger.info(f"{self.service_name} launch_local")
-        asyncio.create_task(self.processing_loop())
+        return asyncio.create_task(self.processing_loop())
 
     # ---- Server based methods ----
 

llama_agents/services/tool.py

@@ -1,11 +1,11 @@
 import asyncio
-import logging
 import uuid
 import uvicorn
 from asyncio import Lock
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
 from pydantic import PrivateAttr
+from logging import getLogger
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from llama_index.core.agent.function_calling.step import (
@@ -29,9 +29,7 @@
     ServiceDefinition,
 )
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-logging.basicConfig(level=logging.DEBUG)
+logger = getLogger(__name__)
 
 
 class ToolService(BaseService):
@@ -198,8 +196,8 @@ def as_consumer(self, remote: bool = False) -> BaseMessageQueueConsumer:
             handler=self.process_message,
         )
 
-    async def launch_local(self) -> None:
-        asyncio.create_task(self.processing_loop())
+    async def launch_local(self) -> asyncio.Task:
+        return asyncio.create_task(self.processing_loop())
 
     # ---- Server based methods ----
 

llama_agents/tools/__init__.py

@@ -1,4 +1,6 @@
 from llama_agents.tools.meta_service_tool import MetaServiceTool
+from llama_agents.tools.service_tool import ServiceTool
+from llama_agents.tools.service_component import ServiceComponent
 
 
-__all__ = ["MetaServiceTool"]
+__all__ = ["MetaServiceTool", "ServiceTool", "ServiceComponent"]