Skip to content

Commit

Permalink
add client
Browse files Browse the repository at this point in the history
  • Loading branch information
nerdai committed Jul 25, 2024
1 parent e23ae7b commit 07f373b
Showing 1 changed file with 52 additions and 2 deletions.
54 changes: 52 additions & 2 deletions snowflake_cybersyn_demo/apps/streamlit.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,70 @@
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]:
for chunk in llama_index_stream:
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")

Expand All @@ -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")
Expand All @@ -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:
Expand Down

0 comments on commit 07f373b

Please sign in to comment.