properly implementation of gradio, langgraph, freeplay and zep (including proper sessions);adding user headshots

api/scripts/create_play_profile_pic.py

@@ -68,10 +68,30 @@ Your output should describe only the image to be generated. No text, captions, o
 #         print("--> skip")
 #     # break
 
+
+image_prompt = """
+A realistic studio-style headshot portrait of {profile[name]}, a {profile[age]}-year-old professional soccer player from {profile[nationality]}. He plays as a {profile[position]} and currently wears jersey number {profile[shirt_number]} for the {team_name} soccer team. 
+
+He is wearing the team's official kit: a {shirt_color} jersey, clean and well-fitted, featuring his number {profile[shirt_number]} on the front. The image should capture a calm and confident expression, suitable for official media. The player has a well-groomed appearance, facing forward with even lighting and a neutral or softly blurred backdrop in studio conditions. No action pose, just a clear, professional profile suitable for player bio pages or trading cards.
+"""
+
+shirt_colors = {
+    "Everglade FC": "emerald green",
+    "Fraser Valley United": "dark red",
+    "Yucatan Force": "dark orange",
+    "Tierra Alta FC": "spring green",
+}
+
 # Generate images for all players
 for player in Player.get_players():
     print(player.name)
-    print(player.profile_pic)
-    response = generate_image(player.profile_pic)
+    # print(player.player_info())
+    shirt_color = shirt_colors[player.team]
+    text = image_prompt.format(
+        profile=player.player_info(),
+        team_name=player.team,
+        shirt_color=shirt_color)
+    # print(text)
+    # print(player.profile_pic)
+    response = generate_image(text)
     player.save_image(response)
-    break

api/scripts/freeplay_playground.py

@@ -1,27 +1,73 @@
 import asyncio
 from pprint import pprint
-from utils.freeplay_helpers import get_formatted_prompt, get_prompt
+from utils.freeplay_helpers import FreeplayClient
 import datetime
+import time
 
-async def main():
-    template = "casual_fan_prompt"
-    environment = "latest"
-    variables = {
-        "now": datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d'),
-        "zep_context": "Zep is your daddy",
-    }
-    history = [
-        {"role": "user", "content": "what are some dinner ideas..."},
-        {"role": "assistant", "content": "here are some dinner ideas..."},
-    ]
+from tools import (
+    PlayerSearchTool,
+    GameSearchTool,
+)
+
+
+available_tools = [
+    GameSearchTool(),
+    PlayerSearchTool(),
+]
+
+
+template = "casual_fan_prompt"
+environment = "latest"
+variables = {
+    "now": datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d'),
+    "zep_context": "Zep is your daddy",
+}
+# history = [
+#     {"role": "user", "content": "what are some dinner ideas..."},
+#     {"role": "assistant", "content": "all you should eat is steamed broccoli - you will live forever"},
+#     {"role": "user", "content": "i don't want to live forever if all i eat is steamed broccoli"},
+#     {"role": "assistant", "content": "silly human"},
+# ]
+# history = [
+#     {'role': 'human', 'content': 'tell me about some players in everglade fc'},
+#     # {'role': 'tool', 'content': '{"number": 23, "name": "Brian Davis", "age": 27, "nationality": "USA", "shirt_number"...c image full of energy and anticipation."}'},
+#     # {'role': 'tool', 'content': '{"number": 10, "name": "Matthew Martin", "age": 24, "nationality": "USA", "shirt_numb...he dramatic flair he brings to the game."}'},
+#     {'role': 'ai', 'content': 'Everglade FC is bursting with talent! Here are some standout players to watch:\n\n1. **...rglade FC brings to the field! Go, team! 🌟'},
+# ]
+history = [
+    {'role': 'user', 'content': 'tell me about some players in everglade fc'},
+    # {'role': 'assistant', 'content': ''},
+    {'role': 'tool_call', 'content': '{"number": 23, "name": "Brian Davis", "age": 27, "nationality": "USA", "shirt_number"...c image full of energy and anticipation."}'},
+    # {'role': 'tool', 'content': '{"number": 10, "name": "Matthew Martin", "age": 24, "nationality": "USA", "shirt_numb...he dramatic flair he brings to the game."}'},
+    {'role': 'assistant', 'content': 'Everglade FC is bursting with talent! Here are some standout players to watch:\n\n1. **...rglade FC brings to the field! Go, team! 🌟'},
+]
 
-    prompt = get_formatted_prompt(template, environment, variables, history=history)
-    # print(prompt.system_content)
-    pprint(prompt.llm_prompt)
+fp_client = FreeplayClient()
+prompt = fp_client.get_prompt(template, environment)
+
+formatted_prompt = prompt.bind(variables, history=history).format()
+
+
+def main():
+    print(formatted_prompt.llm_prompt)
+    print('================')
+    tool_schema = [tool.input_schema.schema_json() for tool in available_tools]
+    print(tool_schema)
+    print('================')
+
+    state = {
+        "start_time": time.time(),
+        "messages": history,
+    }
+    # variables['history'] = history
 
-    # template_prompt = await get_prompt(template, environment)
-    # print(template_prompt.messages)
+    fp_client.create_session()
+    fp_client.record_session(
+        state=state,
+        prompt_vars=variables,
+        formatted_prompt=formatted_prompt,
+    )
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()
     

api/scripts/workflow_playground.py

@@ -2,24 +2,46 @@ import asyncio
 from workflows.base import build_workflow_with_state
 from event_handlers import PrintEventHandler
 from langchain_core.messages import HumanMessage
+from utils.freeplay_helpers import FreeplayClient
+from utils.zep_helpers import ZepClient
 
 from prompts import (
     casual_fan_prompt,
     HumanMessage,
     AIMessage,
 )
+user_id = "[email protected]"
 
 
-workflow, state = build_workflow_with_state(
+# workflow, state = build_workflow_with_state(
+#     handler=PrintEventHandler(),
+#     session_id='5aed14ff09fb415ba77439409f458909',
+#     messages=[
+#         HumanMessage(content="tell me about some players in everglade fc"),
+#     ],
+# )
+
+zep_session_id = ZepClient() \
+                .get_or_create_user(user_id, "Hugh", "Bigly") \
+                .create_session() \
+                .session_id
+freeplay_session_id = FreeplayClient().create_session().session_id
+
+workflow_bundle, state = build_workflow_with_state(
     handler=PrintEventHandler(),
-    session_id='5aed14ff09fb415ba77439409f458909',
+    zep_session_id=zep_session_id,
+    freeplay_session_id=freeplay_session_id,
+    email=user_id,
+    first_name="Hugh",
+    last_name="Bigly",
+    persona="Casual Fan",
     messages=[
         HumanMessage(content="tell me about some players in everglade fc"),
     ],
 )
 
 async def main():
-    await workflow.ainvoke(state)
+    await workflow_bundle.workflow.ainvoke(state)
 
 if __name__ == "__main__":
     asyncio.run(main())

api/scripts/zep_playground2.py

@@ -0,0 +1,19 @@
+import asyncio
+from utils.zep_helpers import ZepClient
+
+zep_client = ZepClient()
+user_id = '[email protected]'
+
+session_id = zep_client\
+    .get_or_create_user(user_id, "Hugh", "Bigly") \
+    .create_session() \
+    .session_id
+print(session_id)
+
+async def main():
+    memory = await zep_client.get_memory_async(session_id)
+    print(f'Memory: {memory}')
+    print(memory.context)
+
+if __name__ == "__main__":
+    asyncio.run(main())

api/server_gradio.py

@@ -6,7 +6,9 @@ from pydantic import BaseModel
 from threading import Thread
 from langchain_core.messages import HumanMessage, AIMessage
 from event_handlers.gradio_handler import GradioEventHandler
-from workflows.base import build_workflow
+from workflows.base import build_workflow_with_state
+from utils.freeplay_helpers import FreeplayClient
+from utils.zep_helpers import ZepClient
 
 lorem_ipsum = """Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."""
 show_state = True
@@ -14,40 +16,98 @@ dev_mode = os.getenv("DEV_MODE", "").lower() == "true"
 
 
 class AppState(BaseModel):
+    email: str = "[email protected]"
+    first_name: str = "Hugh"
+    last_name: str = "Bigly"
     count: int = 0
     persona: str = "Casual Fan"
-    user_query: str = ""
     history: list = []
+    zep_session_id: str = ""
+    freeplay_session_id: str = ""
     
-    def clear(self) -> None:
-        """Reset the state to default values."""
-        self.count = 0
-        self.persona = "Casual Fan"
-        self.user_query = ""
-        self.history = []
+    def ensure_sessions(self):
+        if not self.zep_session_id:
+            self.zep_session_id = ZepClient() \
+                .get_or_create_user(self.email, self.first_name, self.last_name) \
+                .create_session() \
+                .session_id
+        if not self.freeplay_session_id:
+            self.freeplay_session_id = FreeplayClient().create_session().session_id
 
+### Helpers ###
+
+def submit_helper(state, handler, user_query):
+    state.count += 1
+    state.ensure_sessions()
+    message = HumanMessage(content=user_query)
+    state.history.append(message)
+    state = AppState(**state.dict())
+    yield state, ""
+
+    # result = ""
+    # for i in range(0, len(lorem_ipsum), 4):
+    #     time.sleep(0.1)
+    #     result += lorem_ipsum[i:i+4]
+    #     yield state, result
+
+    def start_async_loop():
+        workflow_bundle, workflow_state = build_workflow_with_state(
+            handler=handler,
+            zep_session_id=state.zep_session_id,
+            freeplay_session_id=state.freeplay_session_id,
+            email=state.email,
+            first_name=state.first_name,
+            last_name=state.last_name,
+            persona=state.persona,
+            messages=state.history,
+        )
+
+        async def run_workflow():
+            await workflow_bundle.workflow.ainvoke(workflow_state)
+        
+        asyncio.run(run_workflow())
+
+    thread = Thread(target=start_async_loop, daemon=True)
+    thread.start()
+
+    result = ""
+    while True:
+        token = handler.queue.get()
+        # from colorama import Fore, Style
+        # print(f'{Fore.GREEN}{token}{Style.RESET_ALL}')
+        if token is None:
+            break
+        if isinstance(token, dict):
+            if token["type"] == "info":
+                gr.Info(token["message"])
+                continue
+        result += token
+        yield state, result
+
+    state.history.append(AIMessage(content=result))
+
+### Interface ###
 
 with gr.Blocks() as demo:
     state = gr.State(AppState())
     handler = GradioEventHandler()
-    workflow_state = {
-        "session_id": "5aed14ff09fb415ba77439409f458909",
-        "messages": [],
-    }
 
     gr.Markdown("# Huge League Soccer")
+    with gr.Row(scale=1):
+        email = gr.Textbox(label="Email",
+                           type="email",
+                           interactive=True,
+                           lines=1,
+                           value=state.value.email)
+        first_name = gr.Textbox(label="First Name",
+                                lines=1,
+                                interactive=True,
+                                value=state.value.first_name)
+        last_name = gr.Textbox(label="Last Name",
+                               lines=1,
+                               interactive=True,
+                               value=state.value.last_name)
     llm_response = gr.Textbox(label="LLM Response", lines=10)
-    count = gr.Number(
-        label="Count",
-        interactive=False,
-        visible=show_state,
-    )
-    persona_disp = gr.Text(
-        label="Persona",
-        interactive=False,
-        value=state.value.persona,
-        visible=show_state,
-    )
 
     with gr.Row(scale=1):
         with gr.Column(scale=1):
@@ -67,69 +127,173 @@ with gr.Blocks() as demo:
             submit_btn = gr.Button("Submit", variant="primary", scale=0)
 
     clear_state_btn = gr.Button("Clear State", variant="stop", scale=1)
+
+    with gr.Row(visible=show_state, variant="compact", scale=1):
+        count_disp = gr.Number(
+            label="Count",
+            interactive=False,
+            value=state.value.count,
+            scale=0,
+        )
+        persona_disp = gr.Text(
+            label="Persona",
+            interactive=False,
+            value=state.value.persona,
+            scale=0
+        )
+        zep_session_id_disp = gr.Text(
+            label="Zep Session ID",
+            interactive=False,
+            value=state.value.zep_session_id,
+            scale=1
+        )
+        freeplay_session_id_disp = gr.Text(
+            label="Freeplay Session ID",
+            interactive=False,
+            value=state.value.freeplay_session_id,
+            scale=1
+        )
+
+    ### Events
+
+    @state.change(inputs=[state], outputs=[count_disp, persona_disp, zep_session_id_disp, freeplay_session_id_disp])
+    def state_change(state):
+        return state.count, state.persona, state.zep_session_id, state.freeplay_session_id
+    
+    @clear_state_btn.click(outputs=[state, llm_response, persona, user_query, email, first_name, last_name])
+    def clear_state():
+        new_state = AppState()
+        return new_state, "", new_state.persona, "", new_state.email, new_state.first_name, new_state.last_name
+
+    # @submit_btn.click(inputs=[state, user_query], outputs=[state, llm_response])
+    # def submit(state, user_query):
+    #     state.count += 1
+    #     state.ensure_sessions()
+    #     state = AppState(**state.dict())
+
+    #     user_query = user_query or "tell me about some players in everglade fc"
+    #     message = HumanMessage(content=user_query)
+    #     state.history.append(message)
+
+    #     result = ""
+    #     yield state, result
+    #     for i in range(0, len(lorem_ipsum), 4):
+    #         time.sleep(0.1)
+    #         result += lorem_ipsum[i:i+4]
+    #         yield state, result
+
+
+    @submit_btn.click(inputs=[state, user_query], outputs=[state, llm_response])
+    def submit(state, user_query):
+        user_query = user_query or "tell me about some players in everglade fc"
+        yield from submit_helper(state, handler, user_query)
+        
     
-    @submit_btn.click(inputs=[user_query, state], outputs=[llm_response])
-    def submit(user_query, state):
-        state.user_query = user_query or "tell me about some players in everglade fc"
-        message = HumanMessage(content=state.user_query)
-        state.history.append(message)
-        workflow_state["messages"] = state.history
-
-        def start_async_loop():
-            workflow = build_workflow(handler)
-            async def run_workflow():
-                print(workflow_state)
-                await workflow.ainvoke(workflow_state)
-            asyncio.run(run_workflow())
-
-        thread = Thread(target=start_async_loop, daemon=True)
-        thread.start()
-
-        result = ""
-        while True:
-            token = handler.queue.get()
-            from colorama import Fore, Style
-            print(f'{Fore.GREEN}{token}{Style.RESET_ALL}')
-            if token is None:
-                break
-            if isinstance(token, dict):
-                if token["type"] == "info":
-                    gr.Info(token["message"])
-                    continue
-            result += token
-            yield result
-
-        state.history.append(AIMessage(content=result))
-
-    @user_query.submit(inputs=[user_query, state], outputs=[llm_response])
-    def user_query_change(user_query, state):
-        state.user_query = user_query
-
-        result = state.user_query 
-        for i in range(0, len(lorem_ipsum), 4):
-            time.sleep(0.1)
-            result += lorem_ipsum[i:i+4]
-            yield result
+    # @submit_btn.click(inputs=[user_query, state], outputs=[llm_response, zep_session_id_disp, freeplay_session_id_disp])
+    # def submit(user_query, state):
+    #     state.ensure_sessions()
+    #     state.user_query = user_query or "tell me about some players in everglade fc"
+    #     message = HumanMessage(content=state.user_query)
+    #     state.history.append(message)
+    #     # workflow_state["messages"] = state.history
+
+    #     def start_async_loop():
+    #         workflow, workflow_state = build_workflow_with_state(
+    #             handler=handler,
+    #             zep_session_id=state.zep_session_id,
+    #             freeplay_session_id=state.freeplay_session_id,
+    #             email=state.email,
+    #             first_name=state.first_name,
+    #             last_name=state.last_name,
+    #             persona=state.persona,
+    #             messages=state.history,
+    #         )
+
+    #         async def run_workflow():
+    #             await workflow.ainvoke(workflow_state)
+            
+    #         asyncio.run(run_workflow())
+
+    #     thread = Thread(target=start_async_loop, daemon=True)
+    #     thread.start()
+
+    #     result = ""
+    #     while True:
+    #         token = handler.queue.get()
+    #         # from colorama import Fore, Style
+    #         # print(f'{Fore.GREEN}{token}{Style.RESET_ALL}')
+    #         if token is None:
+    #             break
+    #         if isinstance(token, dict):
+    #             if token["type"] == "info":
+    #                 gr.Info(token["message"])
+    #                 continue
+    #         result += token
+    #         yield result, state.zep_session_id, state.freeplay_session_id
+
+    #     state.history.append(AIMessage(content=result))
+
+    # @user_query.submit(inputs=[user_query, state], outputs=[llm_response, state])
+    # def user_query_change(user_query, state):
+    #     state.user_query = user_query
+    #     state.zep_session_id = f"zep_{int(time.time())}"
+    #     state.freeplay_session_id = f"freeplay_{int(time.time())}"
+
+    #     result = state.user_query 
+    #     for i in range(0, len(lorem_ipsum), 4):
+    #         time.sleep(0.1)
+    #         result += lorem_ipsum[i:i+4]
+    #         yield result, state
 
     @persona.change(inputs=[persona, state], outputs=[persona_disp])
     def persona_change(persona, state):
         state.persona = persona
         return persona
 
-    @submit_btn.click(inputs=[state], outputs=[count])
-    def submit(state):
-        state.count += 1
-        return state.count
+    # @submit_btn.click(inputs=[state], outputs=[state, llm_response])
+    # def submit(state):
+    #     state.count += 1
+    #     new_state = AppState(**state.dict()) 
+    #     result = 'hello'
+    #     for i in range(0, len(lorem_ipsum), 4):
+    #         time.sleep(0.1)
+    #         result += lorem_ipsum[i:i+4]
+    #         yield new_state, result
+
+    # @user_query.submit(inputs=[state], outputs=[count_disp, user_query])
+    # def user_query_change(state):
+    #     state.count += 1
+    #     return state.count, ''
+
+    # @clear_state_btn.click(inputs=[state], outputs=[count_disp, persona_disp, user_query, llm_response, zep_session_id_disp, freeplay_session_id_disp])
+    # def clear_state(state):
+    #     state.clear()
+    #     return state.count, state.persona, state.user_query, "", state.zep_session_id, state.freeplay_session_id
+
+    # @email.change(inputs=[email, state])
+    # def email_change(email, state):
+    #     state.email = email
+
+    # @first_name.change(inputs=[first_name, state])
+    # def first_name_change(first_name, state):
+    #     state.first_name = first_name
+
+    # @last_name.change(inputs=[last_name, state])
+    # def last_name_change(last_name, state):
+    #     state.last_name = last_name
 
-    @user_query.submit(inputs=[state], outputs=[count, user_query])
-    def user_query_change(state):
-        state.count += 1
-        return state.count, ''
+    # @state.change(inputs=[state], outputs=[gr.ParamViewer()])
+    # def state_change(state):
+    #     # return state.zep_session_id, state.freeplay_session_id
+    #     docs = {
+    #         "zep_session_id_ZZZZ": {
+    #             "default": "None\n",
+    #             "type": "str | None\n",
+    #             "description": "Zep session ID.",
+    #         },
+    #     }
 
-    @clear_state_btn.click(inputs=[state], outputs=[count, persona_disp, user_query, llm_response])
-    def clear_state(state):
-        state.clear()
-        return state.count, state.persona, state.user_query, ""
+    #     return gr.ParamViewer(docs)
 
 
 if __name__ == "__main__":

api/tools/__init__.py

@@ -1,7 +1,9 @@
+from langchain_core.documents import Document
 from .player_search import PlayerSearchTool
 from .game_search import GameSearchTool
 
 __all__ = [
     "PlayerSearchTool",
     "GameSearchTool",
+    "Document",
 ]

api/utils/freeplay_helpers.py

@@ -1,58 +1,143 @@
 # import datetime
 import os
+import time
 from functools import lru_cache
 from typing import Union, Optional, List
 from freeplay import Freeplay, RecordPayload, ResponseInfo, CallInfo
+from freeplay.resources.prompts import FormattedPrompt
 from langchain_core.messages import BaseMessage
 
 FREEPLAY_PROJECT_ID = os.getenv("FREEPLAY_PROJECT_ID")
 
 
 # @lru_cache(maxsize=1)
-def get_fp_client():
+def _get_fp_client():
     return Freeplay(
         freeplay_api_key=os.getenv("FREEPLAY_API_KEY"),
         api_base=os.getenv("FREEPLAY_URL"),
     )
 
 
-# retreive and format your prompt
-def get_formatted_prompt(
-    template: str,
-    environment: str = "latest",
-    variables: dict = {},
-    history: Optional[List[BaseMessage]] = None,
-):
-    """
-    Get a formatted prompt from Freeplay.
-    """
-    # get fp client
-    fpClient = get_fp_client()
-    # variables = {**variables, "now": datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d')}
-    # get formatted prompt
-    formatted_prompt = fpClient.prompts.get_formatted(
-        project_id=FREEPLAY_PROJECT_ID,
-        template_name=template,
-        environment=environment,
-        variables=variables,
-        history=history
-    )
-    return formatted_prompt
-
-
-async def get_prompt(
-    template: str,
-    environment: str = "latest",
-):
-    """
-    Get an unformatted prompt template from Freeplay.
-    """
-    # get fp client
-    fpClient = get_fp_client()
-    # get unformatted prompt template
-    template_prompt = fpClient.prompts.get(
-        project_id=FREEPLAY_PROJECT_ID,
-        template_name=template,
-        environment=environment
-    )
-    return template_prompt
+class FreeplayClient:
+
+    def __init__(self, fp_client: Freeplay = None):
+        self.fp_client = fp_client or _get_fp_client()
+        self.session = None
+        self.session_id = None
+        # cache variables for recording
+        self._prompt_cache = {}
+        self._prompt_vars = None
+        self._formatted_prompt = None
+
+    def create_session(self):
+        # create a Freeplay session
+        self.session = self.fp_client.sessions.create()
+        self.session_id = self.session.session_id
+        return self
+
+    @staticmethod
+    def get_fp_client():
+        return _get_fp_client()
+
+    # retreive and format your prompt
+    def get_formatted_prompt(
+        self,
+        template: str,
+        environment: str = "latest",
+        variables: dict = {},
+        history: Optional[List[BaseMessage]] = None,
+    ):
+        """
+        Get a formatted prompt from Freeplay.
+        """
+        formatted_prompt = self.fp_client.prompts.get_formatted(
+            project_id=FREEPLAY_PROJECT_ID,
+            template_name=template,
+            environment=environment,
+            variables=variables,
+            history=history,
+        )
+        return formatted_prompt
+
+    def get_prompt(
+        self,
+        template: str,
+        environment: str = "latest",
+    ):
+        """
+        Get an unformatted prompt template from Freeplay.
+        """
+        key = self._make_cache_key(template, environment)
+        if key in self._prompt_cache:
+            return self._prompt_cache[key]
+        template_prompt = self.fp_client.prompts.get(
+            project_id=FREEPLAY_PROJECT_ID,
+            template_name=template,
+            environment=environment
+        )
+        self._prompt_cache[key] = template_prompt
+        return template_prompt
+
+    def _make_cache_key(self, template: str, environment: str) -> tuple:
+        """
+        Create a cache key for the prompt cache.
+        The key is a tuple of (template, environment).
+        Args:
+            template (str): The prompt template name.
+            environment (str): The environment name.
+        Returns:
+            tuple: (template, environment)
+        """
+        return (template, environment)
+
+    def record_session(
+        self,
+        state,
+        end: Optional[float] = time.time(),
+        formatted_prompt: Optional[FormattedPrompt] = None,
+        prompt_vars: Optional[dict] = None,
+    ):
+        prompt_vars = prompt_vars or self._prompt_vars
+        formatted_prompt = formatted_prompt or self._formatted_prompt
+
+        # convert messages to Freeplay format
+        if state['messages'] and isinstance(state['messages'][0], dict):
+            all_messages = state['messages']
+        else:
+            all_messages = [{'role': m.type, 'content': m.content} for m in state['messages'] if m.content]
+
+        # fix session if we landed here and it's missing
+        if not self.session:
+            self.session = self.fp_client.sessions.restore_session(session_id=state['freeplay_session_id'])
+            self.session_id = self.session.session_id
+
+        # record your LLM call with Freeplay
+        payload = RecordPayload(
+            all_messages=all_messages,
+            inputs=prompt_vars,
+            session_info=self.session, 
+            prompt_info=formatted_prompt.prompt_info,
+            call_info=CallInfo.from_prompt_info(formatted_prompt.prompt_info, start_time=state['start_time'], end_time=end), 
+            response_info=ResponseInfo(
+                # is_complete=chat_response.choices[0].finish_reason == 'stop'
+                is_complete=True
+            )
+        )
+        self.fp_client.recordings.create(payload)
+
+    def get_prompt_by_persona(self,
+                              persona: str,
+                              variables: dict = {},
+                              history: Optional[List[BaseMessage]] = None):
+        if 'casual' in persona.lower():
+            prompt = self.get_prompt(template='casual_fan_prompt', environment='latest')
+        elif 'super' in persona.lower():
+            prompt = self.get_prompt(template='super_fan_prompt', environment='latest')
+        else:
+            raise ValueError(f"Unknown persona: {persona}")
+
+        formatted_prompt = prompt.bind(variables=variables, history=history).format()
+        self._prompt_vars = variables
+        self._formatted_prompt = formatted_prompt
+
+        return formatted_prompt

api/utils/zep_helpers.py

@@ -1,30 +1,133 @@
+import uuid
 from functools import lru_cache
 from langchain_core.messages import BaseMessage
-from zep_cloud.client import AsyncZep
+from zep_cloud.client import AsyncZep, Zep
 from zep_cloud.types import Message
+from zep_cloud.errors import NotFoundError
+from typing import Optional
 
 
 # @lru_cache(maxsize=1)
-def get_zep_client():
+def _get_zep_client():
     return AsyncZep()
 
 
-async def record_session(session_id: str, messages: list[BaseMessage]):
-    if len(messages) >= 2:
-        user_message = messages[0]
-        assistant_message = messages[-1]
-        messages = [
+class ZepClient:
+
+    def __init__(self):
+        self.zep_client_async = AsyncZep()
+        self.zep_client = Zep()
+        self._user = None
+        self.session_id = None
+
+    async def get_memory_async(self, session_id: Optional[str] = None):
+        session_id = session_id or self.session_id
+        if session_id is None:
+            raise ValueError("No session ID provided")
+        return await self.zep_client_async.memory.get(session_id=session_id)
+
+    # async def get_session(self, session_id: str):
+    #     return await self.zep_client.session.get(session_id=session_id)
+
+    # async def get_or_create_user(self, email: str, first_name: Optional[str] = None, last_name: Optional[str] = None):
+    #     if self._user and self._user.user_id == email:
+    #         return self
+    #     try:
+    #         self._user = await self.get_user(email)
+    #     except NotFoundError:
+    #         self._user = await self.create_user(email, first_name, last_name)
+    #     return self
+
+    def get_or_create_user(self, email: str, first_name: Optional[str] = None, last_name: Optional[str] = None):
+        if self._user and self._user.user_id == email:
+            return self
+        try:
+            self.get_user(email)
+        except NotFoundError:
+            self.create_user(email, first_name, last_name)
+        return self
+
+    def create_session(self):
+        session_id = uuid.uuid4().hex # A new session identifier
+        self.zep_client.memory.add_session(
+            session_id=session_id,
+            user_id=self._user.user_id,
+        )
+        self.session_id = session_id
+        return self
+
+    def create_user(self, email: str, first_name: Optional[str] = None, last_name: Optional[str] = None):
+        self._user = self.zep_client.user.add(
+            user_id=email,
+            email=email,
+            first_name=first_name,
+            last_name=last_name,
+        )
+        return self
+
+    # async def create_user(self, email: str, first_name: Optional[str] = None, last_name: Optional[str] = None):
+    #     self._user = await self.zep_client.user.add(
+    #         user_id=email,
+    #         email=email,
+    #         first_name=first_name,
+    #         last_name=last_name,
+    #     )
+    #     return self._user
+
+    # async def get_user(self, email: str):
+    #     self._user = await self.zep_client.user.get(user_id=email)
+    #     return self._user
+
+    def get_user(self, email: str):
+        self._user = self.zep_client.user.get(user_id=email)
+        return self
+
+    @staticmethod
+    def get_zep_client() -> AsyncZep:
+        return _get_zep_client()
+
+    async def record_session(self,
+                             messages: list[BaseMessage],
+                             session_id: Optional[str] = None):
+        session_id = session_id or self.session_id
+        if session_id is None:
+            raise ValueError("No session ID provided")
+        if len(messages) >= 2:
+            user_message = messages[0]
+            assistant_message = messages[-1]
+            messages = [
+                Message(
+                    role="user",
+                content=user_message.content,
+                role_type="user",
+            ),
             Message(
-                role="user",
-            content=user_message.content,
-            role_type="user",
-        ),
-        Message(
-            role="assistant",
-            content=assistant_message.content,
-            role_type="assistant",
-        ),
-    ]
-    zep_client = get_zep_client()
-    await zep_client.memory.add(session_id=session_id, messages=messages)
+                role="assistant",
+                content=assistant_message.content,
+                role_type="assistant",
+            ),
+        ]
+        await self.zep_client_async.memory.add(session_id=session_id, messages=messages)
+
+
+# async def record_session(session_id: str,
+#                          messages: list[BaseMessage],
+#                          zep_client: AsyncZep = None):
+#     if len(messages) >= 2:
+#         user_message = messages[0]
+#         assistant_message = messages[-1]
+#         messages = [
+#             Message(
+#                 role="user",
+#             content=user_message.content,
+#             role_type="user",
+#         ),
+#         Message(
+#             role="assistant",
+#             content=assistant_message.content,
+#             role_type="assistant",
+#         ),
+#     ]
+#     zep_client = zep_client or get_zep_client()
+#     await zep_client.memory.add(session_id=session_id, messages=messages)
     

api/workflows/base.py

@@ -1,6 +1,9 @@
 import asyncio
 import datetime
 import operator
+import time
+from dataclasses import dataclass
+from typing import Any, Tuple, List, Optional
 from functools import partial
 from typing import TypedDict, Annotated, Sequence
 from langchain_openai import ChatOpenAI
@@ -16,16 +19,13 @@ from langgraph.graph import StateGraph, END
 from langchain_core.callbacks import AsyncCallbackHandler
 from langgraph.prebuilt import ToolNode
 
-
-from utils.zep_helpers import (
-    AsyncZep,
-    get_zep_client,
-    record_session,
-)
+from utils.freeplay_helpers import FreeplayClient
+from utils.zep_helpers import ZepClient 
 from prompts import (
     casual_fan_prompt,
 )
 from tools import (
+    Document,
     PlayerSearchTool,
     GameSearchTool,
 )
@@ -42,29 +42,53 @@ llm_with_tools = llm.bind_tools(tools=available_tools)
 
 
 class AgentState(TypedDict):
-    session_id: str
+    email: str
+    first_name: str
+    last_name: str
+    zep_session_id: str
+    freeplay_session_id: str
+    persona: str
+    start_time: float
+    zep_memory: Optional[str] = None
     messages: Annotated[Sequence[BaseMessage], operator.add]
 
 
-async def call_model(state: AgentState, handler: AsyncCallbackHandler, zep_client: AsyncZep) -> dict:
-    session_id = state["session_id"]
-    memory = await zep_client.memory.get(session_id=session_id)
+async def call_model(state: AgentState,
+                     handler: AsyncCallbackHandler,
+                     zep_client: ZepClient,
+                     freeplay_client: FreeplayClient) -> dict:
+    zep_session_id = state["zep_session_id"]
     messages = state["messages"]
+    persona = state["persona"]
     
-    prompt = casual_fan_prompt.format(
-        zep_context=memory.context,
-        input=messages,
-        now=datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d'),
+    memory = state["zep_memory"]
+    if memory is None:
+        memory = await zep_client.get_memory_async(session_id=zep_session_id)
+        memory = memory.context or "New user/session. No memory."
+    
+    # old_prompt = casual_fan_prompt.format(
+    #     zep_context=memory,
+    #     input=messages,
+    #     now=datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d'),
+    # )
+    prompt = freeplay_client.get_prompt_by_persona(
+        persona=persona,
+        variables={
+            "now": datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%d'),
+            "zep_context": memory,
+        },
+        history=messages,
     )
 
     # response = await llm_with_tools.ainvoke(prompt, stream=True,
     #                                     config={"callbacks" :[handler]})
-    response = await llm_with_tools.with_config(callbacks=[handler]).ainvoke(prompt, stream=True)
+    response = await llm_with_tools.with_config(callbacks=[handler]).ainvoke(prompt.llm_prompt, stream=True)
 
-    return {'messages': [response]}
+    return {'messages': [response], 'zep_memory': memory}
 
 
-async def call_tool(state: AgentState, handler: AsyncCallbackHandler) -> dict:
+async def call_tool(state: AgentState,
+                    handler: AsyncCallbackHandler) -> dict:
     messages = state["messages"]
     last_message = messages[-1]
     tools_by_name = {tool.name: tool for tool in available_tools}
@@ -81,56 +105,113 @@ async def call_tool(state: AgentState, handler: AsyncCallbackHandler) -> dict:
 
     results = []
     for observation, tool_call in zip(observations, last_message.tool_calls):
+        if isinstance(observation, list) and len(observation) > 0 and isinstance(observation[0], Document):
+            observation = "\n\n".join(doc.page_content for doc in observation)
         results.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
     return {"messages": results}
 
 
-async def should_continue(state: AgentState, handler: AsyncCallbackHandler, zep_client: AsyncZep) -> str:
+async def should_continue(state: AgentState,
+                          handler: AsyncCallbackHandler,
+                          zep_client: ZepClient,
+                          freeplay_client: FreeplayClient) -> str:
     messages = state["messages"]
     last_message = messages[-1]
     if 'tool_calls' not in last_message.additional_kwargs:
         # inform zep of final response
-        await record_session(state["session_id"], messages)
+        await zep_client.record_session(
+            session_id=state["zep_session_id"],
+            messages=messages,
+        )
+        # inform freeplay of final response
+        # TODO: wait for feedback from freeplay about tool messages
+        # freeplay_client.record_session(state)
+        # trigger on_workflow_end callback
         if hasattr(handler, 'on_workflow_end'):
             await handler.on_workflow_end(state)
         return 'end'
     return 'continue'
 
+### WorkflowBundle ###
+@dataclass
+class WorkflowBundle:
+    """
+    Bundle containing the compiled workflow and its core dependencies.
+
+    Attributes:
+        workflow: The compiled workflow object ready for execution.
+        zep_client: The Zep client instance used for memory/context.
+        freeplay_client: The Freeplay client instance for session tracking.
+    """
+    workflow: Any
+    zep_client: ZepClient
+    freeplay_client: FreeplayClient
+
 
-def build_workflow(handler: AsyncCallbackHandler):
+def build_workflow(handler: AsyncCallbackHandler) -> WorkflowBundle:
+    """
+    Build and compile the workflow, returning all core dependencies as a WorkflowBundle.
+
+    Args:
+        handler: AsyncCallbackHandler for this workflow.
+
+    Returns:
+        WorkflowBundle: Contains compiled workflow, zep_client, and freeplay_client.
+
+    Example usage:
+        bundle = build_workflow(handler)
+        # Access workflow: bundle.workflow
+        # Access zep client: bundle.zep_client
+        # Access freeplay client: bundle.freeplay_client
+    """
     workflow = StateGraph(AgentState)
-    zep_client = get_zep_client()
-    workflow.add_node('agent', partial(call_model, handler=handler, zep_client=zep_client))
+    zep_client = ZepClient()
+    freeplay_client = FreeplayClient()
+    workflow.add_node('agent', partial(call_model, handler=handler, zep_client=zep_client, freeplay_client=freeplay_client))
     workflow.add_node('tools', partial(call_tool, handler=handler))
     workflow.set_entry_point('agent')
 
     workflow.add_conditional_edges(
         'agent',
-        partial(should_continue, handler=handler, zep_client=zep_client),
+        partial(should_continue, handler=handler, zep_client=zep_client, freeplay_client=freeplay_client),
         {
             'continue': 'tools',
             'end': END,
-        }   
+        }
     )
 
     workflow.add_edge('tools', 'agent')
 
-    return workflow.compile()
+    return WorkflowBundle(workflow=workflow.compile(), zep_client=zep_client, freeplay_client=freeplay_client)
 
 
-def build_workflow_with_state(handler: AsyncCallbackHandler, session_id: str, messages=None):
+def build_workflow_with_state(handler: AsyncCallbackHandler,
+                              zep_session_id: Optional[str] = None,
+                              freeplay_session_id: Optional[str] = None,
+                              email: Optional[str] = None,
+                              first_name: Optional[str] = None,
+                              last_name: Optional[str] = None,
+                              persona: Optional[str] = None,
+                              messages: Optional[List[BaseMessage]] = None) -> Tuple[WorkflowBundle, AgentState]:
     """
     Utility to build workflow and initial state in one step.
     Args:
         handler: AsyncCallbackHandler for this workflow
-        session_id: unique session identifier
+        zep_session_id: unique session identifier
         messages: optional initial message list
     Returns:
         (workflow, state) tuple ready for execution
     """
-    workflow = build_workflow(handler)
+    bundle = build_workflow(handler)
     state = {
-        "session_id": session_id,
+        "zep_session_id": zep_session_id,
+        "freeplay_session_id": freeplay_session_id,
+        "email": email,
+        "first_name": first_name,
+        "last_name": last_name,
+        "persona": persona,
         "messages": messages or [],
+        "start_time": time.time(),
+        "zep_memory": None,
     }
-    return workflow, state
+    return bundle, state