First version of AI Trainer assistant

app.py

@@ -1,64 +1,257 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import os
+import time
+from tools import create_workout_table_graph
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+class WorkoutTableGUI:
+    def __init__(self, graph=None, share=False):
+        self.graph = graph if graph else create_workout_table_graph()
+        self.share = share
+        self.partial_message = ""
+        self.response = {}
+        self.max_iterations = 10
+        self.iterations = []
+        self.threads = []
+        self.thread_id = -1
+        self.thread = {"configurable": {"thread_id": str(self.thread_id)}}
+        self.demo = self.create_interface()
 
+    def run_agent(self, start, topic, stop_after):
+        if start:
+            self.iterations.append(0)
+            config = {'task': topic, "max_revisions": 2, "revision_number": 0,
+                      'lnode': "", 'planner': "no plan", 'draft': "no draft", 'critique': "no critique", 
+                      'content': ["no content",], 'queries': "no queries", 'count': 0}
+            self.thread_id += 1  # new agent, new thread
+            self.threads.append(self.thread_id)
+        else:
+            config = None
+        self.thread = {"configurable": {"thread_id": str(self.thread_id)}}
+        while self.iterations[self.thread_id] < self.max_iterations:
+            self.response = self.graph.invoke(config, self.thread)
+            self.iterations[self.thread_id] += 1
+            self.partial_message += str(self.response)
+            self.partial_message += f"\n------------------\n\n"
+            lnode, nnode, _, rev, acount = self.get_disp_state()
+            yield self.partial_message, lnode, nnode, self.thread_id, rev, acount
+            config = None
+            if not nnode:  
+                return
+            if lnode in stop_after:
+                return
+        return
+    
+    def get_disp_state(self):
+        current_state = self.graph.get_state(self.thread)
+        lnode = current_state.values["lnode"]
+        acount = current_state.values["count"]
+        rev = current_state.values["revision_number"]
+        nnode = current_state.next
+        return lnode, nnode, self.thread_id, rev, acount
+    
+    def get_state(self, key):
+        current_values = self.graph.get_state(self.thread)
+        if key in current_values.values:
+            lnode, nnode, self.thread_id, rev, astep = self.get_disp_state()
+            new_label = f"last_node: {lnode}, thread_id: {self.thread_id}, rev: {rev}, step: {astep}"
+            return gr.update(label=new_label, value=current_values.values[key])
+        else:
+            return ""  
+    
+    def get_content(self):
+        current_values = self.graph.get_state(self.thread)
+        if "content" in current_values.values:
+            content = current_values.values["content"]
+            lnode, nnode, thread_id, rev, astep = self.get_disp_state()
+            new_label = f"last_node: {lnode}, thread_id: {self.thread_id}, rev: {rev}, step: {astep}"
+            return gr.update(label=new_label, value="\n\n".join(item for item in content) + "\n\n")
+        else:
+            return ""  
+    
+    def update_hist_pd(self):
+        hist = []
+        for state in self.graph.get_state_history(self.thread):
+            if state.metadata['step'] < 1:
+                continue
+            thread_ts = state.config['configurable']['thread_ts']
+            tid = state.config['configurable']['thread_id']
+            count = state.values['count']
+            lnode = state.values['lnode']
+            rev = state.values['revision_number']
+            nnode = state.next
+            st = f"{tid}:{count}:{lnode}:{nnode}:{rev}:{thread_ts}"
+            hist.append(st)
+        return gr.Dropdown(label="update_state from: thread:count:last_node:next_node:rev:thread_ts", 
+                           choices=hist, value=hist[0] if hist else None, interactive=True)
+    
+    def find_config(self, thread_ts):
+        for state in self.graph.get_state_history(self.thread):
+            config = state.config
+            if config['configurable']['thread_ts'] == thread_ts:
+                return config
+        return None
+            
+    def copy_state(self, hist_str):
+        thread_ts = hist_str.split(":")[-1]
+        config = self.find_config(thread_ts)
+        state = self.graph.get_state(config)
+        self.graph.update_state(self.thread, state.values, as_node=state.values['lnode'])
+        new_state = self.graph.get_state(self.thread)
+        new_thread_ts = new_state.config['configurable']['thread_ts']
+        tid = new_state.config['configurable']['thread_id']
+        count = new_state.values['count']
+        lnode = new_state.values['lnode']
+        rev = new_state.values['revision_number']
+        nnode = new_state.next
+        return lnode, nnode, new_thread_ts, rev, count
+    
+    def update_thread_pd(self):
+        return gr.Dropdown(label="choose thread", choices=self.threads, value=self.thread_id, interactive=True)
+    
+    def switch_thread(self, new_thread_id):
+        self.thread = {"configurable": {"thread_id": str(new_thread_id)}}
+        self.thread_id = new_thread_id
+        return 
+    
+    def modify_state(self, key, asnode, new_state):
+        current_values = self.graph.get_state(self.thread)
+        current_values.values[key] = new_state
+        self.graph.update_state(self.thread, current_values.values, as_node=asnode)
+        return
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+    def create_interface(self):
+        with gr.Blocks(theme=gr.themes.Default(spacing_size='sm', text_size="sm")) as demo:
+            
+            def updt_disp():
+                current_state = self.graph.get_state(self.thread)
+                hist = []
+                for state in self.graph.get_state_history(self.thread):
+                    if state.metadata['step'] < 1:
+                        continue
+                    s_thread_ts = state.config['configurable']['thread_ts']
+                    s_tid = state.config['configurable']['thread_id']
+                    s_count = state.values['count']
+                    s_lnode = state.values['lnode']
+                    s_rev = state.values['revision_number']
+                    s_nnode = state.next
+                    st = f"{s_tid}:{s_count}:{s_lnode}:{s_nnode}:{s_rev}:{s_thread_ts}"
+                    hist.append(st)
+                if not current_state.metadata:
+                    return {}
+                else:
+                    return {
+                        topic_bx: current_state.values["task"],
+                        lnode_bx: current_state.values["lnode"],
+                        count_bx: current_state.values["count"],
+                        revision_bx: current_state.values["revision_number"],
+                        nnode_bx: current_state.next,
+                        threadid_bx: self.thread_id,
+                        thread_pd: gr.Dropdown(label="choose thread", choices=self.threads, value=self.thread_id, interactive=True),
+                        step_pd: gr.Dropdown(label="update_state from: thread:count:last_node:next_node:rev:thread_ts", 
+                               choices=hist, value=hist[0] if hist else None, interactive=True),
+                    }
+            
+            def get_snapshots():
+                new_label = f"thread_id: {self.thread_id}, Summary of snapshots"
+                sstate = ""
+                for state in self.graph.get_state_history(self.thread):
+                    for key in ['plan', 'draft', 'critique']:
+                        if key in state.values:
+                            state.values[key] = state.values[key][:80] + "..."
+                    if 'content' in state.values:
+                        for i in range(len(state.values['content'])):
+                            state.values['content'][i] = state.values['content'][i][:20] + '...'
+                    if 'writes' in state.metadata:
+                        state.metadata['writes'] = "not shown"
+                    sstate += str(state) + "\n\n"
+                return gr.update(label=new_label, value=sstate)
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+            def vary_btn(stat):
+                return gr.update(variant=stat)
+            
+            with gr.Tab("Agent"):
+                with gr.Row():
+                    topic_bx = gr.Textbox(label="Workout Table", value="Workout Table for a 30 year old male who wants to gain muscle mass and strength")
+                    gen_btn = gr.Button("Generate Workout Table", scale=0, min_width=80, variant='primary')
+                    cont_btn = gr.Button("Continue Workout Table", scale=0, min_width=80)
+                with gr.Row():
+                    lnode_bx = gr.Textbox(label="last node", min_width=100)
+                    nnode_bx = gr.Textbox(label="next node", min_width=100)
+                    threadid_bx = gr.Textbox(label="Thread", scale=0, min_width=80)
+                    revision_bx = gr.Textbox(label="Draft Rev", scale=0, min_width=80)
+                    count_bx = gr.Textbox(label="count", scale=0, min_width=80)
+                with gr.Accordion("Manage Agent", open=False):
+                    checks = list(self.graph.nodes.keys())
+                    checks.remove('__start__')
+                    stop_after = gr.CheckboxGroup(checks, label="Interrupt After State", value=checks, scale=0, min_width=400)
+                    with gr.Row():
+                        thread_pd = gr.Dropdown(choices=self.threads, interactive=True, label="select thread", min_width=120, scale=0)
+                        step_pd = gr.Dropdown(choices=['N/A'], interactive=True, label="select step", min_width=160, scale=1)
+                live = gr.Textbox(label="Live Agent Output", lines=5, max_lines=5)
+        
+                # actions
+                sdisps = [topic_bx, lnode_bx, nnode_bx, threadid_bx, revision_bx, count_bx, step_pd, thread_pd]
+                thread_pd.input(self.switch_thread, [thread_pd], None).then(
+                                fn=updt_disp, inputs=None, outputs=sdisps)
+                step_pd.input(self.copy_state, [step_pd], None).then(
+                              fn=updt_disp, inputs=None, outputs=sdisps)
+                gen_btn.click(vary_btn, gr.Number("secondary", visible=False), gen_btn).then(
+                              fn=self.run_agent, inputs=[gr.Number(True, visible=False), topic_bx, stop_after], outputs=[live], show_progress=True).then(
+                              fn=updt_disp, inputs=None, outputs=sdisps).then( 
+                              vary_btn, gr.Number("primary", visible=False), gen_btn).then(
+                              vary_btn, gr.Number("primary", visible=False), cont_btn)
+                cont_btn.click(vary_btn, gr.Number("secondary", visible=False), cont_btn).then(
+                               fn=self.run_agent, inputs=[gr.Number(False, visible=False), topic_bx, stop_after], 
+                               outputs=[live]).then(
+                               fn=updt_disp, inputs=None, outputs=sdisps).then(
+                               vary_btn, gr.Number("primary", visible=False), cont_btn)
+        
+            with gr.Tab("Workout Table"):
+                with gr.Row():
+                    refresh_btn = gr.Button("Refresh")
+                    modify_btn = gr.Button("Modify")
+                plan = gr.Textbox(label="Plan", lines=10, interactive=True)
+                refresh_btn.click(fn=self.get_state, inputs=gr.Number("plan", visible=False), outputs=plan)
+                modify_btn.click(fn=self.modify_state, inputs=[gr.Number("plan", visible=False),
+                                                          gr.Number("planner", visible=False), plan], outputs=None).then(
+                                 fn=updt_disp, inputs=None, outputs=sdisps)
+            with gr.Tab("Research Content"):
+                refresh_btn = gr.Button("Refresh")
+                content_bx = gr.Textbox(label="content", lines=10)
+                refresh_btn.click(fn=self.get_content, inputs=None, outputs=content_bx)
+            with gr.Tab("Draft"):
+                with gr.Row():
+                    refresh_btn = gr.Button("Refresh")
+                    modify_btn = gr.Button("Modify")
+                draft_bx = gr.Textbox(label="draft", lines=10, interactive=True)
+                refresh_btn.click(fn=self.get_state, inputs=gr.Number("draft", visible=False), outputs=draft_bx)
+                modify_btn.click(fn=self.modify_state, inputs=[gr.Number("draft", visible=False),
+                                                          gr.Number("generate", visible=False), draft_bx], outputs=None).then(
+                                fn=updt_disp, inputs=None, outputs=sdisps)
+            with gr.Tab("Feedback"):
+                with gr.Row():
+                    refresh_btn = gr.Button("Refresh")
+                    modify_btn = gr.Button("Modify")
+                feedback_bx = gr.Textbox(label="Feedback", lines=10, interactive=True)
+                refresh_btn.click(fn=self.get_state, inputs=gr.Number("feedback", visible=False), outputs=feedback_bx)
+                modify_btn.click(fn=self.modify_state, inputs=[gr.Number("feedback", visible=False),
+                                                          gr.Number("reflect", visible=False), 
+                                                          feedback_bx], outputs=None).then(
+                                fn=updt_disp, inputs=None, outputs=sdisps)
+            with gr.Tab("StateSnapShots"):
+                with gr.Row():
+                    refresh_btn = gr.Button("Refresh")
+                snapshots = gr.Textbox(label="State Snapshots Summaries")
+                refresh_btn.click(fn=get_snapshots, inputs=None, outputs=snapshots)
+        return demo
 
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+    def launch(self, share=None):
+        if port := os.getenv("PORT1"):
+            self.demo.launch(share=True, server_port=int(port), server_name="0.0.0.0")
+        else:
+            self.demo.launch(share=self.share)
 
 
 if __name__ == "__main__":
-    demo.launch()
+    workout_table = WorkoutTableGUI()
+    workout_table.launch()

helper.py

@@ -0,0 +1,22 @@
+import warnings
+warnings.filterwarnings("ignore", message=".*TqdmWarning.*")
+from dotenv import load_dotenv
+
+_ = load_dotenv()
+
+from tools import create_workout_table_graph
+from app import WorkoutTableGUI
+
+# This file is kept for backward compatibility
+# The functionality has been moved to tools.py and app.py
+
+def create_workout_table():
+    """
+    Create an workout table instance with the graph and GUI
+    """
+    graph = create_workout_table_graph()
+    return WorkoutTableGUI(graph)
+
+if __name__ == "__main__":
+    workout_table = create_workout_table()
+    workout_table.launch()

tools.py

@@ -0,0 +1,202 @@
+import warnings
+warnings.filterwarnings("ignore", message=".*TqdmWarning.*")
+from dotenv import load_dotenv
+
+_ = load_dotenv()
+
+from langgraph.graph import StateGraph, END
+from typing import TypedDict, Annotated, List
+import operator
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage, ChatMessage
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel
+from tavily import TavilyClient
+import os
+import sqlite3
+
+class AgentState(TypedDict):
+    task: str
+    lnode: str
+    plan: str
+    draft: str
+    feedback: str
+    content: List[str]
+    queries: List[str]
+    revision_number: int
+    max_revisions: int
+    count: Annotated[int, operator.add]
+
+
+class Queries(BaseModel):
+    queries: List[str]
+
+# Tool functions
+def plan_node(model, state: AgentState):
+    table_output = """
+    Workout Table Sequence:
+    - Section 1: Warm-up
+    - Section 2: Strength Training
+    - Section 3: Cardio
+    - Section 4: Cool-down
+
+    Workout Table Example:
+    Workout Table Full Body (3 times per week):
+    Day 1:
+    Squat on Chair: 2-3 sets of 8-10 repetitions
+    Bench Press with Dumbbells: 2-3 sets of 8-10 repetitions
+    Row on the machine: 2-3 sets of 8-10 repetitions
+    Glute Machine or Glute Bridge: 2-3 sets of 10-12 repetitions
+    Development with dumbbells: 2-3 sets of 8-10 repetitions
+    Dumbbell curls: 2-3 sets of 10-12 repetitions
+
+    Day 2:
+    Leg Extension: 2-3 sets of 10-12 repetitions
+    Arms pushdowns (supporting on knees): 2-3 sets of 8-10 repetitions
+    Dumbbell row: 2-3 sets of 8-10 repetitions
+    Calf Raises (standing): 2-3 sets of 10-12 repetitions
+    Lateral Raises: 2-3 sets of 10-12 repetitions
+    Triceps Extension: 2-3 sets of 10-12 repetitions
+
+    Day 3:
+    Cardio (bike or treadmill): 30 minutes at moderate pace
+    Abdominal: 2-3 sets of 12-15 repetitions
+    Plank: 2-3 sets, holding for 20-30 seconds
+
+    Remember to always warm up before training and stretch after. 
+    """
+
+    PLAN_PROMPT = ("You are an expert gym trainer tasked with writing a high level workout table. "
+                    "Write such an outline for the user provided workout. Give the three main headers of an outline of "
+                     "the workout table along with any relevant notes or instructions for the sections. "
+                     f"Here is the user's workout table: {table_output}")
+    messages = [
+        SystemMessage(content=PLAN_PROMPT), 
+        HumanMessage(content=state['task'])
+    ]
+    response = model.invoke(messages)
+    return {"plan": response.content,
+           "lnode": "planner",
+            "count": 1,
+           }
+
+def research_plan_node(model, tavily, state: AgentState):
+    RESEARCH_PLAN_PROMPT = ("You are a researcher charged with providing information that can "
+                             "be used when writing the following workout table. Generate a list of search "
+                             "queries that will gather "
+                             "any relevant information. Only generate 3 queries max.")
+    queries = model.with_structured_output(Queries).invoke([
+        SystemMessage(content=RESEARCH_PLAN_PROMPT),
+        HumanMessage(content=state['task'])
+    ])
+    content = state['content'] or []  # add to content
+    for q in queries.queries:
+        response = tavily.search(query=q, max_results=2)
+        for r in response['results']:
+            content.append(r['content'])
+    return {"content": content,
+            "queries": queries.queries,
+           "lnode": "research_plan",
+            "count": 1,
+           }
+
+def generation_node(model, state: AgentState):
+    WRITER_PROMPT = ("You are an gym trainer assistant tasked with writing excellent workout tables. "
+                      "Generate the best workout table possible for the user's request and the initial outline. "
+                      "If the user provides feedback, respond with a revised version of your previous attempts. "
+                      "Utilize all the information below as needed: \n"
+                      "------\n"
+                      "{content}")
+    content = "\n\n".join(state['content'] or [])
+    user_message = HumanMessage(
+        content=f"{state['task']}\n\nHere is my workout table:\n\n{state['plan']}")
+    messages = [
+        SystemMessage(
+            content=WRITER_PROMPT.format(content=content)
+        ),
+        user_message
+        ]
+    response = model.invoke(messages)
+    return {
+        "draft": response.content, 
+        "revision_number": state.get("revision_number", 1) + 1,
+        "lnode": "generate",
+        "count": 1,
+    }
+
+def reflection_node(model, state: AgentState):
+    REFLECTION_PROMPT = ("You are an instructor personal grading an workout table submission. "
+                          "Generate feedback and recommendations for the user's submission. "
+                          "Provide detailed recommendations, including requests for objectives, level of intensity, health benefits, health conditions, etc.")
+    messages = [
+        SystemMessage(content=REFLECTION_PROMPT), 
+        HumanMessage(content=state['draft'])
+    ]
+    response = model.invoke(messages)
+    return {"feedback": response.content,
+           "lnode": "reflect",
+            "count": 1,
+    }
+
+def research_feedback_node(model, tavily, state: AgentState):
+    RESEARCH_FEEDBACK_PROMPT = ("You are a researcher charged with providing information that can "
+                "be used when writing the following workout table. Generate a list of search "
+                "queries that will gather "
+                "any relevant information. Only generate 3 queries max.")
+    queries = model.with_structured_output(Queries).invoke([
+        SystemMessage(content=RESEARCH_FEEDBACK_PROMPT),
+        HumanMessage(content=state['feedback'])
+    ])
+    content = state['content'] or []  # add to content
+    for q in queries.queries:
+        response = tavily.search(query=q, max_results=2)
+        for r in response['results']:
+            content.append(r['content'])
+    queries = model.with_structured_output(Queries).invoke([
+        SystemMessage(content=RESEARCH_FEEDBACK_PROMPT),
+        HumanMessage(content=state['feedback'])
+    ])
+    content = state['content'] or []
+    for q in queries.queries:
+        response = tavily.search(query=q, max_results=2)
+        for r in response['results']:
+            content.append(r['content'])
+    return {"content": content,
+           "lnode": "research_feedback",
+            "count": 1,
+    }
+
+def should_continue(state):
+    if state["revision_number"] > state["max_revisions"]:
+        return END
+    return "reflect"
+
+# Function to create the graph
+def create_workout_table_graph():
+    model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+    tavily = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+    
+    builder = StateGraph(AgentState)
+    builder.add_node("planner", lambda state: plan_node(model, state))
+    builder.add_node("research_plan", lambda state: research_plan_node(model, tavily, state))
+    builder.add_node("generate", lambda state: generation_node(model, state))
+    builder.add_node("reflect", lambda state: reflection_node(model, state))
+    builder.add_node("research_feedback", lambda state: research_feedback_node(model, tavily, state))
+    builder.set_entry_point("planner")
+    builder.add_conditional_edges(
+        "generate", 
+        should_continue, 
+        {END: END, "reflect": "reflect"}
+    )
+    builder.add_edge("planner", "research_plan")
+    builder.add_edge("research_plan", "generate")
+    builder.add_edge("reflect", "research_feedback")
+    builder.add_edge("research_feedback", "generate")
+    memory = SqliteSaver(conn=sqlite3.connect(":memory:", check_same_thread=False))
+    graph = builder.compile(
+        checkpointer=memory,
+        interrupt_after=['planner', 'generate', 'reflect', 'research_plan', 'research_feedback']
+    )
+    
+    return graph
+