polojuan/agentic-workflows

1 2import json3import re4from datetime import datetime5import streamlit as st6from .research_agent import research_agent7from .editor_agent import editor_agent8from .writer_agent import writer_agent9from .medical_agent import medical_agent10import time11from openai import OpenAI12from dotenv import find_dotenv, load_dotenv13 14# Load environment variables15load_dotenv(find_dotenv())16 17 18def agent_register(page):19    if page == "researcher":20        agent_registry = {21            "research_agent": research_agent,22            "editor_agent": editor_agent,23            "writer_agent": writer_agent,24        }25    elif page == "medical":26        agent_registry = {27            "medical_agent": medical_agent,28            "editor_agent": editor_agent,29            "writer_agent": writer_agent,30        }31    return agent_registry32 33def clean_json_block(raw: str) -> str:34    """35    Clean the contents of a JSON block that may come wrapped with Markdown backticks.36    """37    raw = raw.strip()38    if raw.startswith("```"):39        raw = re.sub(r"^```(?:json)?\n?", "", raw)40        raw = re.sub(r"\n?```$", "", raw)41    return raw.strip()42 43 44def executor_agent(plan_steps: list[str], model: str = "gpt-5-mini", page: str = "researcher"):45    # Get client from session state46    client = st.session_state.get("client") or OpenAI()47 48    history = []49 50    print("==================================")51    print("🎯 Execution Agent")52    print("==================================")53    if "steps" not in st.session_state:54        st.session_state.steps = [container1, container2, container3, container4, container5] = [None]*555        st.session_state.expanders = st.expander("Agent Steps", expanded=True)56 57    total_used_token = 058    with st.session_state.expanders:59        for i, step in enumerate(plan_steps):60            st.session_state.steps[i] = st.container(border=True)61            st.session_state.steps[i].write(f"Step {i+1}: {step}")62    63    for i, step in enumerate(plan_steps):64        agent_decision_prompt = f"""65        You are an execution manager for a multi-agent research team.66        67        Given the following instruction, identify which agent should perform it and extract the clean task.68        69        Return only a valid JSON object with two keys:70        - "agent": one of {list(agent_register(page).keys())}71        - "task": a string with the instruction that the agent should follow72        73        Only respond with a valid JSON object. Do not include explanations or markdown formatting.74        75        Instruction: "{step}"76        """77        response = client.chat.completions.create(78            model=model,79            messages=[{"role": "user", "content": agent_decision_prompt}]80            )81        82        raw_content = response.choices[0].message.content83        cleaned_json = clean_json_block(raw_content)84        agent_info = json.loads(cleaned_json)85        86        agent_name = agent_info["agent"]87        task = agent_info["task"]88 89        context = "\n".join([90            f"Step {j+1} executed by {a}:\n{r}" 91            for j, (s, a, r) in enumerate(history)92        ])93        enriched_task = f"""You are {agent_name}.94        95        Here is the context of what has been done so far:96        {context}97        98        Your next task is:99        {task}100        """101 102        print(f"\n🛠️ Executing with agent: `{agent_name}` on task: {task}")103        agent_registry = agent_register(page)104        if agent_name in agent_registry:105            with st.session_state.steps[i]:106                start_time = time.time()107                with st.spinner(f"Executing... ", show_time=True):108                    output, used_token = agent_registry[agent_name](enriched_task, model=st.session_state.model)109                    history.append((step, agent_name, output))110                    total_used_token += used_token111                    print(f"✅ Agent Used Tokens:\n{used_token}")112                    elapsed_time = time.time() - start_time113                    print(f"✅ Elapsed Time: {elapsed_time:.2f} seconds")114                    st.success(f"✅ Completed with {used_token} token used in {elapsed_time:.2f} seconds!")115        else:116            with st.session_state.steps[i]:117                start_time = time.time()118                with st.spinner(f"Executing... ", show_time=True):119                    output, used_token = f"⚠️ Unknown agent: {agent_name}"120                    history.append((step, agent_name, output))121                    total_used_token += used_token122                    print(f"✅ Agent Used Tokens:\n{used_token}")123                    elapsed_time = time.time() - start_time124                    print(f"✅ Elapsed Time: {elapsed_time:.2f} seconds")125                    st.success(f"✅ Completed with {used_token} token used in {elapsed_time:.2f} seconds!")126            127    print(f"✅ Output:\n{output}")128    print(f"✅ Total Tokens Used:\n{total_used_token}")129        130    return history, total_used_token131 132