Initial commit from Create Llama

Author: marcusschiesser

.devcontainer/devcontainer.json

@@ -0,0 +1,47 @@
+{
+  "image": "mcr.microsoft.com/vscode/devcontainers/typescript-node:dev-20-bullseye",
+  "features": {
+    "ghcr.io/devcontainers-contrib/features/turborepo-npm:1": {},
+    "ghcr.io/devcontainers-contrib/features/typescript:2": {},
+    "ghcr.io/devcontainers/features/python:1": {
+      "version": "3.11",
+      "toolsToInstall": [
+        "flake8",
+        "black",
+        "mypy",
+        "poetry"
+      ]
+    }
+  },
+  "customizations": {
+    "codespaces": {
+      "openFiles": [
+        "README.md"
+      ]
+    },
+    "vscode": {
+      "extensions": [
+        "ms-vscode.typescript-language-features",
+        "esbenp.prettier-vscode",
+        "ms-python.python",
+        "ms-python.black-formatter",
+        "ms-python.vscode-flake8",
+        "ms-python.vscode-pylance"
+      ],
+      "settings": {
+        "python.formatting.provider": "black",
+        "python.languageServer": "Pylance",
+        "python.analysis.typeCheckingMode": "basic"
+      }
+    }
+  },
+  "containerEnv": {
+    "POETRY_VIRTUALENVS_CREATE": "false",
+    "PYTHONPATH": "${PYTHONPATH}:${workspaceFolder}"
+  },
+  "forwardPorts": [
+    3000,
+    8000
+  ],
+  "postCreateCommand": "poetry install"
+}

.env

@@ -0,0 +1,57 @@
+# The Llama Cloud API key.
+# LLAMA_CLOUD_API_KEY=
+
+# The provider for the AI models to use.
+MODEL_PROVIDER=openai
+
+# The name of LLM model to use.
+MODEL=gpt-4o-mini
+
+# Name of the embedding model to use.
+EMBEDDING_MODEL=text-embedding-3-large
+
+# Dimension of the embedding model to use.
+EMBEDDING_DIM=1024
+
+# The questions to help users get started (multi-line).
+# CONVERSATION_STARTERS=
+
+# The OpenAI API key to use.
+OPENAI_API_KEY=sk-proj-7qNCi2Btib0Y7AADy3tgT3BlbkFJ9d7IFV2lk2IfRnvDQbxp
+
+# Temperature for sampling from the model.
+# LLM_TEMPERATURE=
+
+# Maximum number of tokens to generate.
+# LLM_MAX_TOKENS=
+
+# The number of similar embeddings to return when retrieving documents.
+TOP_K=3
+
+# The time in milliseconds to wait for the stream to return a response.
+STREAM_TIMEOUT=60000
+
+# FILESERVER_URL_PREFIX is the URL prefix of the server storing the images generated by the interpreter.
+FILESERVER_URL_PREFIX=http://localhost:8000/api/files
+
+# The address to start the backend app.
+APP_HOST=0.0.0.0
+
+# The port to start the backend app.
+APP_PORT=8000
+
+# MESSAGE_QUEUE_PORT=
+
+# CONTROL_PLANE_PORT=
+
+# HUMAN_CONSUMER_PORT=
+
+AGENT_QUERY_ENGINE_PORT=8003
+
+AGENT_QUERY_ENGINE_DESCRIPTION=Query information from the provided data
+
+AGENT_DUMMY_PORT=8004
+
+# The system prompt for the AI model.
+SYSTEM_PROMPT=You are a helpful assistant who helps users with their questions.
+

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.11 as build
+
+WORKDIR /app
+
+ENV PYTHONPATH=/app
+
+# Install Poetry
+RUN curl -sSL https://install.python-poetry.org | POETRY_HOME=/opt/poetry python && \
+    cd /usr/local/bin && \
+    ln -s /opt/poetry/bin/poetry && \
+    poetry config virtualenvs.create false
+
+# Install Chromium for web loader
+# Can disable this if you don't use the web loader to reduce the image size
+RUN apt update && apt install -y chromium chromium-driver
+
+# Install dependencies
+COPY ./pyproject.toml ./poetry.lock* /app/
+RUN poetry install --no-root --no-cache --only main
+
+# ====================================
+FROM build as release
+
+COPY . .
+
+CMD ["python", "main.py"]

README.md

@@ -0,0 +1,50 @@
+This is a [LlamaIndex](https://www.llamaindex.ai/) project using [FastAPI](https://fastapi.tiangolo.com/) bootstrapped with [`create-llama`](https://github.com/run-llama/LlamaIndexTS/tree/main/packages/create-llama).
+
+## Getting Started
+
+First, setup the environment with poetry:
+
+> **_Note:_** This step is not needed if you are using the dev-container.
+
+```shell
+poetry install
+poetry shell
+```
+
+Then check the parameters that have been pre-configured in the `.env` file in this directory. (E.g. you might need to configure an `OPENAI_API_KEY` if you're using OpenAI as model provider).
+
+Second, generate the embeddings of the documents in the `./data` directory (if this folder exists - otherwise, skip this step):
+
+```shell
+poetry run generate
+```
+
+Third, run all the services in one command:
+
+```shell
+poetry run python main.py
+```
+
+You can monitor and test the agent services with `llama-agents` monitor TUI:
+
+```shell
+poetry run llama-agents monitor --control-plane-url http://127.0.0.1:8001
+```
+
+## Services:
+
+- Message queue (port 8000): To exchange the message between services
+- Control plane (port 8001): A gateway to manage the tasks and services.
+- Human consumer (port 8002): To handle result when the task is completed.
+- Agent service `query_engine` (port 8003): Agent that can query information from the configured LlamaIndex index.
+- Agent service `dummy_agent` (port 8004): A dummy agent that does nothing. Good starting point to add more agents.
+
+The ports listed above are set by default, but you can change them in the `.env` file.
+
+## Learn More
+
+To learn more about LlamaIndex, take a look at the following resources:
+
+- [LlamaIndex Documentation](https://docs.llamaindex.ai) - learn about LlamaIndex.
+
+You can check out [the LlamaIndex GitHub repository](https://github.com/run-llama/llama_index) - your feedback and contributions are welcome!

app/agents/dummy/agent.py

@@ -0,0 +1,33 @@
+from llama_agents import AgentService, SimpleMessageQueue
+from llama_index.core.agent import FunctionCallingAgentWorker
+from llama_index.core.tools import FunctionTool
+from llama_index.core.settings import Settings
+from app.utils import load_from_env
+
+
+DEFAULT_DUMMY_AGENT_DESCRIPTION = "I'm a dummy agent which does nothing."
+
+
+def dummy_function():
+    """
+    This function does nothing.
+    """
+    return ""
+
+
+def init_dummy_agent(message_queue: SimpleMessageQueue) -> AgentService:
+    agent = FunctionCallingAgentWorker(
+        tools=[FunctionTool.from_defaults(fn=dummy_function)],
+        llm=Settings.llm,
+        prefix_messages=[],
+    ).as_agent()
+
+    return AgentService(
+        service_name="dummy_agent",
+        agent=agent,
+        message_queue=message_queue.client,
+        description=load_from_env("AGENT_DUMMY_DESCRIPTION", throw_error=False)
+        or DEFAULT_DUMMY_AGENT_DESCRIPTION,
+        host=load_from_env("AGENT_DUMMY_HOST", throw_error=False) or "127.0.0.1",
+        port=int(load_from_env("AGENT_DUMMY_PORT")),
+    )

app/agents/query_engine/agent.py

@@ -0,0 +1,52 @@
+import os
+from llama_agents import AgentService, SimpleMessageQueue
+from llama_index.core.agent import FunctionCallingAgentWorker
+from llama_index.core.tools import QueryEngineTool, ToolMetadata
+from llama_index.core.settings import Settings
+from app.engine.index import get_index
+from app.utils import load_from_env
+
+
+DEFAULT_QUERY_ENGINE_AGENT_DESCRIPTION = (
+    "Used to answer the questions using the provided context data."
+)
+
+
+def get_query_engine_tool() -> QueryEngineTool:
+    """
+    Provide an agent worker that can be used to query the index.
+    """
+    index = get_index()
+    if index is None:
+        raise ValueError("Index not found. Please create an index first.")
+    query_engine = index.as_query_engine(similarity_top_k=int(os.getenv("TOP_K", 3)))
+    return QueryEngineTool(
+        query_engine=query_engine,
+        metadata=ToolMetadata(
+            name="context_data",
+            description="""
+                Provide the provided context information. 
+                Use a detailed plain text question as input to the tool.
+            """,
+        ),
+    )
+
+
+def init_query_engine_agent(
+    message_queue: SimpleMessageQueue,
+) -> AgentService:
+    """
+    Initialize the agent service.
+    """
+    agent = FunctionCallingAgentWorker(
+        tools=[get_query_engine_tool()], llm=Settings.llm, prefix_messages=[]
+    ).as_agent()
+    return AgentService(
+        service_name="context_query_agent",
+        agent=agent,
+        message_queue=message_queue.client,
+        description=load_from_env("AGENT_QUERY_ENGINE_DESCRIPTION", throw_error=False)
+        or DEFAULT_QUERY_ENGINE_AGENT_DESCRIPTION,
+        host=load_from_env("AGENT_QUERY_ENGINE_HOST", throw_error=False) or "127.0.0.1",
+        port=int(load_from_env("AGENT_QUERY_ENGINE_PORT")),
+    )

app/core/control_plane.py

@@ -0,0 +1,19 @@
+from llama_index.llms.openai import OpenAI
+from llama_agents import AgentOrchestrator, ControlPlaneServer
+from app.core.message_queue import message_queue
+from app.utils import load_from_env
+
+
+control_plane_host = (
+    load_from_env("CONTROL_PLANE_HOST", throw_error=False) or "127.0.0.1"
+)
+control_plane_port = load_from_env("CONTROL_PLANE_PORT", throw_error=False) or "8001"
+
+
+# setup control plane
+control_plane = ControlPlaneServer(
+    message_queue=message_queue,
+    orchestrator=AgentOrchestrator(llm=OpenAI()),
+    host=control_plane_host,
+    port=int(control_plane_port) if control_plane_port else None,
+)

app/core/message_queue.py

@@ -0,0 +1,12 @@
+from llama_agents import SimpleMessageQueue
+from app.utils import load_from_env
+
+message_queue_host = (
+    load_from_env("MESSAGE_QUEUE_HOST", throw_error=False) or "127.0.0.1"
+)
+message_queue_port = load_from_env("MESSAGE_QUEUE_PORT", throw_error=False) or "8000"
+
+message_queue = SimpleMessageQueue(
+    host=message_queue_host,
+    port=int(message_queue_port) if message_queue_port else None,
+)

app/core/task_result.py

@@ -0,0 +1,88 @@
+import json
+from logging import getLogger
+from pathlib import Path
+from fastapi import FastAPI
+from typing import Dict, Optional
+from llama_agents import CallableMessageConsumer, QueueMessage
+from llama_agents.message_queues.base import BaseMessageQueue
+from llama_agents.message_consumers.base import BaseMessageQueueConsumer
+from llama_agents.message_consumers.remote import RemoteMessageConsumer
+from app.utils import load_from_env
+from app.core.message_queue import message_queue
+
+
+logger = getLogger(__name__)
+
+
+class TaskResultService:
+    def __init__(
+        self,
+        message_queue: BaseMessageQueue,
+        name: str = "human",
+        host: str = "127.0.0.1",
+        port: Optional[int] = 8002,
+    ) -> None:
+        self.name = name
+        self.host = host
+        self.port = port
+
+        self._message_queue = message_queue
+
+        # app
+        self._app = FastAPI()
+        self._app.add_api_route(
+            "/", self.home, methods=["GET"], tags=["Human Consumer"]
+        )
+        self._app.add_api_route(
+            "/process_message",
+            self.process_message,
+            methods=["POST"],
+            tags=["Human Consumer"],
+        )
+
+    @property
+    def message_queue(self) -> BaseMessageQueue:
+        return self._message_queue
+
+    def as_consumer(self, remote: bool = False) -> BaseMessageQueueConsumer:
+        if remote:
+            return RemoteMessageConsumer(
+                url=(
+                    f"http://{self.host}:{self.port}/process_message"
+                    if self.port
+                    else f"http://{self.host}/process_message"
+                ),
+                message_type=self.name,
+            )
+
+        return CallableMessageConsumer(
+            message_type=self.name,
+            handler=self.process_message,
+        )
+
+    async def process_message(self, message: QueueMessage) -> None:
+        Path("task_results").mkdir(exist_ok=True)
+        with open("task_results/task_results.json", "+a") as f:
+            json.dump(message.model_dump(), f)
+            f.write("\n")
+
+    async def home(self) -> Dict[str, str]:
+        return {"message": "hello, human."}
+
+    async def register_to_message_queue(self) -> None:
+        """Register to the message queue."""
+        await self.message_queue.register_consumer(self.as_consumer(remote=True))
+
+
+human_consumer_host = (
+    load_from_env("HUMAN_CONSUMER_HOST", throw_error=False) or "127.0.0.1"
+)
+human_consumer_port = load_from_env("HUMAN_CONSUMER_PORT", throw_error=False) or "8002"
+
+
+human_consumer_server = TaskResultService(
+    message_queue=message_queue,
+    host=human_consumer_host,
+    port=int(human_consumer_port) if human_consumer_port else None,
+    name="human",
+)