start outsourcing to an app controller

snowflake_cybersyn_demo/apps/controller.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Generator, List, Optional
+
+import streamlit as st
+from llama_agents import LlamaAgentsClient
+from llama_index.core.llms import ChatMessage, ChatResponseGen
+
+
+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)
+
+
+class Controller:
+    def __init__(self, llama_agents_client: Optional[LlamaAgentsClient]):
+        self._client = llama_agents_client
+
+    def _llama_index_stream_wrapper(
+        self,
+        llama_index_stream: ChatResponseGen,
+    ) -> Generator[str, Any, Any]:
+        for chunk in llama_index_stream:
+            yield chunk.delta
+
+    def _handle_task_submission(
+        self, 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 = self._client.create_task(task_input)
+        # task = TaskModel(
+        #     task_id=task_id,
+        #     input=task_input,
+        #     chat_history=[
+        #         ChatMessage(role="user", content=task_input),
+        #         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)

snowflake_cybersyn_demo/apps/streamlit.py

@@ -1,70 +1,24 @@
-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.core.llms import ChatMessage
 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]:
-    for chunk in llama_index_stream:
-        yield chunk.delta
-
-
-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)
-
+from snowflake_cybersyn_demo.apps.controller import Controller
 
 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}"
-    )
-)
+# 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}"
+#     )
+# )
+llama_agents_client = None
+controller = Controller(llama_agents_client)
+
+### App
 st.set_page_config(layout="wide")
 st.title("Human In The Loop W/ LlamaAgents")
 
@@ -84,7 +38,7 @@ def _handle_task_submission(llama_agents_client: LlamaAgentsClient) -> None:
         "Task input",
         placeholder="Enter a task input.",
         key="task_input",
-        on_change=_handle_task_submission,
+        on_change=controller._handle_task_submission,
         args=(llama_agents_client,),
     )
 
@@ -108,7 +62,9 @@ def _handle_task_submission(llama_agents_client: LlamaAgentsClient) -> None:
                     for m in st.session_state.messages
                 ]
             )
-            response = st.write_stream(_llama_index_stream_wrapper(stream))
+            response = st.write_stream(
+                controller._llama_index_stream_wrapper(stream)
+            )
         st.session_state.messages.append(
             {"role": "assistant", "content": response}
         )