From 6cce05fa27f1fd011cf2424e879af777b1392cd8 Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Wed, 4 Sep 2024 15:23:10 +0700
Subject: [PATCH] feat: switch fast api / main endpoint

---
 .env                    |  3 +++
 README.md               |  9 ++++++---
 app/examples/factory.py |  4 ++--
 main.py                 | 38 +++++++++++++++++++++++++++++++++++++-
 4 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/.env b/.env
index 44b3703..24b77c0 100644
--- a/.env
+++ b/.env
@@ -30,3 +30,6 @@ TOP_K=3
 
 # Choose 'choreography', 'orchestrator', or 'workflow' for the type of agent interaction to use.
 EXAMPLE_TYPE=workflow
+
+# Set it to true to start FastAPI endpoint
+FAST_API=true
\ No newline at end of file
diff --git a/README.md b/README.md
index cba43fa..9c25975 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,11 @@
-This is a [LlamaIndex](https://www.llamaindex.ai/) multi-agents project using [Workflows](https://docs.llamaindex.ai/en/stable/understanding/workflows/). 
+This is a [LlamaIndex](https://www.llamaindex.ai/) multi-agents project using [Workflows](https://docs.llamaindex.ai/en/stable/understanding/workflows/).
 
 ## Overview
 
 This example is using three agents to generate a blog post:
-- a researcher that retrieves content via a RAG pipeline, 
-- a writer that specializes in writing blog posts and 
+
+- a researcher that retrieves content via a RAG pipeline,
+- a writer that specializes in writing blog posts and
 - a reviewer that is reviewing the blog post.
 
 There are three different methods how the agents can interact to reach their goal:
@@ -39,6 +40,8 @@ poetry run python main.py
 
 Per default, the example is using the explicit workflow. You can change the example by setting the `EXAMPLE_TYPE` environment variable to `choreography` or `orchestrator`.
 
+To add an API endpoint, set the `FAST_API` environment variable to `true`.
+
 ## Learn More
 
 To learn more about LlamaIndex, take a look at the following resources:
diff --git a/app/examples/factory.py b/app/examples/factory.py
index d79cd7b..9ae36b6 100644
--- a/app/examples/factory.py
+++ b/app/examples/factory.py
@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 from app.examples.choreography import create_choreography
 from app.examples.orchestrator import create_orchestrator
 from app.examples.workflow import create_workflow
@@ -14,7 +14,7 @@ import os
 logger = logging.getLogger("uvicorn")
 
 
-def create_agent(chat_history: List[ChatMessage]) -> Workflow:
+def create_agent(chat_history: Optional[List[ChatMessage]] = None) -> Workflow:
     agent_type = os.getenv("EXAMPLE_TYPE", "").lower()
     match agent_type:
         case "choreography":
diff --git a/main.py b/main.py
index 11395a0..c3f4668 100644
--- a/main.py
+++ b/main.py
@@ -1,8 +1,12 @@
 # flake8: noqa: E402
+import asyncio
 import os
+import textwrap
+from typing import AsyncGenerator
 from dotenv import load_dotenv
 
 from app.config import DATA_DIR
+from app.examples.factory import create_agent
 
 load_dotenv()
 
@@ -64,9 +68,41 @@ app.include_router(config_router, prefix="/api/chat/config")
 app.include_router(file_upload_router, prefix="/api/chat/upload")
 
 
-if __name__ == "__main__":
+def run_api():
     app_host = os.getenv("APP_HOST", "0.0.0.0")
     app_port = int(os.getenv("APP_PORT", "8000"))
     reload = True if environment == "dev" else False
 
     uvicorn.run(app="main:app", host=app_host, port=app_port, reload=reload)
+
+
+async def main():
+    def info(prefix: str, text: str) -> None:
+        truncated = textwrap.shorten(text, width=255, placeholder="...")
+        print(f"[{prefix}] {truncated}")
+
+    agent = create_agent()
+
+    task = asyncio.create_task(
+        agent.run(
+            input="Write a blog post about physical standards for letters",
+            streaming=True,
+        )
+    )
+
+    async for ev in agent.stream_events():
+        info(ev.name, ev.msg)
+
+    ret: AsyncGenerator = await task
+    async for token in ret:
+        print(token.delta, end="", flush=True)
+
+    # ret: AgentRunResult = await task
+    # print(ret.response.message.content)
+
+
+if __name__ == "__main__":
+    if os.getenv("FAST_API", "false").lower() == "false":
+        asyncio.run(main())
+    else:
+        run_api()