From 07f373ba1516368ed0d3ab9cb9f1378177ea6d0e Mon Sep 17 00:00:00 2001
From: Andrei Fajardo <andrei@nerdai.io>
Date: Thu, 25 Jul 2024 01:58:48 -0400
Subject: [PATCH] add client

---
 snowflake_cybersyn_demo/apps/streamlit.py | 54 ++++++++++++++++++++++-
 1 file changed, 52 insertions(+), 2 deletions(-)

diff --git a/snowflake_cybersyn_demo/apps/streamlit.py b/snowflake_cybersyn_demo/apps/streamlit.py
index fc266a8..584b9c7 100644
--- a/snowflake_cybersyn_demo/apps/streamlit.py
+++ b/snowflake_cybersyn_demo/apps/streamlit.py
@@ -1,11 +1,29 @@
-from typing import Any, Generator
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Generator, List, Optional
 
 import pandas as pd
 import streamlit as st
+from llama_agents import LlamaAgentsClient
 from llama_index.core.llms import ChatMessage, ChatResponseGen
 from llama_index.llms.openai import OpenAI
 
 
+class TaskStatus(str, Enum):
+    HUMAN_REQUIRED = "human_required"
+    COMPLETED = "completed"
+    SUBMITTED = "submitted"
+
+
+@dataclass
+class TaskModel:
+    task_id: str
+    input: str
+    status: TaskStatus
+    prompt: Optional[str] = None
+    chat_history: List[ChatMessage] = field(default_factory=list)
+
+
 def _llama_index_stream_wrapper(
     llama_index_stream: ChatResponseGen,
 ) -> Generator[str, Any, Any]:
@@ -13,11 +31,40 @@ def _llama_index_stream_wrapper(
         yield chunk.delta
 
 
-def _handle_task_submission() -> None:
+def _handle_task_submission(llama_agents_client: LlamaAgentsClient) -> None:
+    """Handle the user submitted message. Clear task submission box, and
+    add the new task to the submitted list.
+    """
+
+    # create new task and store in state
+    task_input = st.session_state.task_input
+    task_id = llama_agents_client.create_task(task_input)
+    task = TaskModel(
+        task_id=task_id,
+        input=task_input,
+        chat_history=[
+            message,
+            ChatMessage(
+                role="assistant",
+                content=f"Successfully submitted task: {task_id}.",
+            ),
+        ],
+        status=TaskStatus.SUBMITTED,
+    )
     st.session_state.submitted_pills.append(st.session_state.task_input)
+    st.session_state.tasks.append(task)
 
 
 llm = OpenAI(model="gpt-4o-mini")
+control_plane_host = "0.0.0.0"
+control_plane_port = 8001
+llama_agents_client = LlamaAgentsClient(
+    control_plane_url=(
+        f"http://{control_plane_host}:{control_plane_port}"
+        if control_plane_port
+        else f"http://{control_plane_host}"
+    )
+)
 st.set_page_config(layout="wide")
 st.title("Human In The Loop W/ LlamaAgents")
 
@@ -26,6 +73,8 @@ def _handle_task_submission() -> None:
     st.session_state["submitted_pills"] = []
 st.session_state["human_required_pills"] = []
 st.session_state["completed_pills"] = []
+if "tasks" not in st.session_state:
+    st.session_state["tasks"] = []
 
 
 left, right = st.columns([1, 2], vertical_alignment="bottom")
@@ -36,6 +85,7 @@ def _handle_task_submission() -> None:
         placeholder="Enter a task input.",
         key="task_input",
         on_change=_handle_task_submission,
+        args=(llama_agents_client,),
     )
 
 with right: