veadk-python 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

veadk/cli/studio/fast_api.py

@@ -11,30 +11,37 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+import json
 import os
+import uuid
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
+from google.adk.agents.run_config import RunConfig, StreamingMode
+from google.genai import types
+from pydantic import BaseModel
 
-from veadk.cli.studio.agent_processor import AgentProcessor
-from veadk.cli.studio.model import (
+from veadk.cli.studio.models import (
     GetAgentResponse,
-    GetAgentsResponse,
     GetEventResponse,
-    GetHistorySessionsResponse,
-    GetMemoryResponse,
-    OptimizePromptRequest,
-    OptimizePromptResponse,
-    ReplacePromptRequest,
-    RunAgentRequest,
-    RunAgentResponse,
-    SetAgentRequest,
-    SetAgentResponse,
-    SetRunnerRequest,
-    TraceAgentResponse,
+    GetHistoryMessagesResponse,
+    GetTracingResponse,
+    Message,
+    RunnerConfig,
+    SpanItem,
+    TextPart,
+    ToolPart,
 )
+from veadk.cli.studio.studio_processor import StudioProcessor
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+processor = None
 
 app = FastAPI()
 
@@ -44,9 +51,6 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-    allow_origin_regex=None,
-    expose_headers=[],
-    max_age=600,
 )
 
 NEXT_STATIC_DIR = os.path.join(os.path.dirname(__file__), "web")
@@ -60,158 +64,307 @@ async def read_root():
         return FileResponse(index_path, media_type="text/html")
 
 
-agent_processor = AgentProcessor()
-
-
-# health check
-@app.get("/ping")
-async def root():
-    return {"message": "pong"}
-
-
-# get all agents
-@app.get("/agents")
-async def get_agents() -> GetAgentsResponse:
-    agents_dir = agent_processor.agents_dir
-    agents = agent_processor.get_agent_names()
-    return GetAgentsResponse(agents_dir=agents_dir, agents=agents)
-
-
-# set current agent
-@app.post("/set_agent")
-async def set_agent(request: SetAgentRequest) -> SetAgentResponse:
-    agent_processor.set_agent(
-        agent_name=request.agent_name,
-    )
-    return SetAgentResponse(
-        name=agent_processor.agent.name,
-        description=agent_processor.agent.description,
-        model_name=agent_processor.agent.model_name,
-        instruction=agent_processor.agent.instruction,
-        long_term_memory_backend=agent_processor.agent.long_term_memory.backend
-        if agent_processor.agent.long_term_memory
-        else "",
-        knowledgebase_backend=agent_processor.agent.knowledgebase.backend
-        if agent_processor.agent.knowledgebase
-        else "",
-    )
-
-
-@app.post("/set_runner")
-async def set_runner(request: SetRunnerRequest):
-    await agent_processor.set_runner(
-        app_name=request.app_name,
-        user_id=request.user_id,
-        session_id=request.session_id,
-        short_term_memory_backend=request.short_term_memory_backend,
-        short_term_memory_db_url=request.short_term_memory_db_url,
+@app.get("/get_history_messages")
+async def get_history_messages() -> GetHistoryMessagesResponse:
+    session_service = processor.short_term_memory.session_service
+    session = await session_service.get_session(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
     )
-    # TODO: return long & short term memory state
-    return {"message": f"Runner set to {request.app_name}"}
 
-
-@app.get("/memory")
-async def get_memory() -> GetMemoryResponse:
-    memory_status = agent_processor.get_memory_status()
-    return GetMemoryResponse(
-        **memory_status,
+    # prevent no session created
+    if session:
+        messages = []
+        tool_mapping = {}
+        for event in session.events:
+            formatted_parts = []
+            for part in event.content.parts:
+                if part.text:
+                    formatted_part = TextPart(text=part.text)
+                    formatted_parts.append(formatted_part)
+                if part.function_call:
+                    formatted_part = ToolPart(
+                        type="tool-" + part.function_call.name,
+                        state="input-available",
+                        toolCallId=part.function_call.id,
+                        input=part.function_call.args,
+                    )
+                    tool_mapping[part.function_call.id] = formatted_part
+                    formatted_parts.append(formatted_part)
+                if part.function_response:
+                    formatted_part = tool_mapping[part.function_response.id]
+                    formatted_part.output = part.function_response.response
+                    formatted_part.state = "output-available"
+
+            if formatted_parts:  # prevent only function_response in event
+                session_message = Message(
+                    id=event.id,
+                    role=event.author if event.author == "user" else "assistant",
+                    parts=formatted_parts,
+                )
+                messages.append(session_message)
+        return GetHistoryMessagesResponse(messages=messages)
+    else:
+        return GetHistoryMessagesResponse(messages=[])
+
+
+async def runner_run_sse(user_text: str):
+    message = types.Content(role="user", parts=[types.Part(text=user_text)])
+
+    await processor.runner.short_term_memory.create_session(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
     )
 
-
-# get current agent
-@app.get("/agent")
-async def get_agent() -> GetAgentResponse:
+    logger.info("Begin to process user message under SSE method.")
+
+    # message begin
+    msg_id = f"msg-{str(uuid.uuid4())}"
+    MESSAGE_START = {"type": "start", "messageId": msg_id}
+    yield f"data: {json.dumps(MESSAGE_START)}\n\n"
+
+    text_id = f"text-{str(uuid.uuid4())}"
+    TEXT_START = {"type": "text-start", "id": text_id}
+    yield f"data: {json.dumps(TEXT_START)}\n\n"
+
+    async for event in processor.runner.runner.run_async(
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
+        new_message=message,
+        run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+    ):
+        for function_call in event.get_function_calls():
+            TOOL_INPUT_START = {
+                "type": "tool-input-start",
+                "toolCallId": function_call.id,
+                "toolName": function_call.name,
+            }
+            yield f"data: {json.dumps(TOOL_INPUT_START)}\n\n"
+
+            TOOL_INPUT_AVAILABLE = {
+                "type": "tool-input-available",
+                "toolCallId": function_call.id,
+                "toolName": function_call.name,
+                "input": function_call.args,
+            }
+            yield f"data: {json.dumps(TOOL_INPUT_AVAILABLE)}\n\n"
+
+        for function_response in event.get_function_responses():
+            # for load_memory and load_knowledgebase tools, the output is based BaseModel
+            # we have to convert it
+            for key, value in function_response.response.items():
+                if isinstance(value, BaseModel):
+                    function_response.response[key] = value.model_dump(
+                        exclude_none=False
+                    )
+
+            TOOL_OUTPUT_AVAILABLE = {
+                "type": "tool-output-available",
+                "toolCallId": function_response.id,
+                "output": function_response.response,
+            }
+            yield f"data: {json.dumps(TOOL_OUTPUT_AVAILABLE)}\n\n"
+
+        if (
+            not event.is_final_response()
+            and len(event.content.parts) > 0
+            and event.content.parts[0].text
+        ):
+            TEXT_DELTA = {
+                "type": "text-delta",
+                "delta": event.content.parts[0].text,
+                "id": text_id,
+            }
+            yield f"data: {json.dumps(TEXT_DELTA)}\n\n"
+
+        if event.is_final_response():
+            TEXT_END = {"type": "text-end", "id": text_id}
+            yield f"data: {json.dumps(TEXT_END)}\n\n"
+
+            TEXT_FINISH = {"type": "finish", "finishReason": "stop"}
+            yield f"data: {json.dumps(TEXT_FINISH)}\n\n"
+
+            DONE = "[DONE]"
+            yield f"data: {DONE}\n\n"
+
+    logger.info("SSE Stream Completed.")
+
+
+@app.post("/run_sse")
+async def run_sse(request: Request):
+    data = await request.json()
+    user_text = data["messages"][-1]["parts"][0]["text"]
+
+    response = StreamingResponse(runner_run_sse(user_text=user_text))
+    response.headers["x-vercel-ai-ui-message-stream"] = "v1"
+    response.headers["Content-Type"] = "text/event-stream"
+    return response
+
+
+@app.get("/get_agent")
+def get_agent() -> GetAgentResponse:
     return GetAgentResponse(
-        name=agent_processor.agent.name,
-        description=agent_processor.agent.description,
-        model_name=agent_processor.agent.model_name,
-        instruction=agent_processor.agent.instruction,
+        name=processor.agent.name,
+        description=processor.agent.description,
+        instruction=processor.agent.instruction,
+        model_name=processor.agent.model_name,
+        short_term_memory_backend="",
+        short_term_memory_db_url="",
+        long_term_memory_backend="",
+        knowledgebase_backend="",
     )
 
 
-@app.post("/optimize_prompt")
-async def optimize_prompt(request: OptimizePromptRequest) -> OptimizePromptResponse:
-    current_prompt = request.prompt
-    feedback = request.feedback
-    prompt = agent_processor.optimize_prompt(current_prompt, feedback)
-    return OptimizePromptResponse(prompt=prompt)
+def parse_tracing_file(filepath: str):
+    with open(filepath, "r") as f:
+        spans = json.load(f)
+
+    max_latency = -1
+    # 1. generate span mapping
+    store: dict[str, SpanItem] = {}
+    for span in spans:
+        event_id = ""
+        # in case of root node
+        if "gcp.vertex.agent.event_id" in span["attributes"]:
+            event_id = str(span["attributes"]["gcp.vertex.agent.event_id"])
+
+        latency = (span["end_time"] - span["start_time"]) / 1000000
+        max_latency = max(max_latency, latency)
+        span_item = SpanItem(
+            name=span["name"],
+            span_id=str(span["span_id"]),
+            trace_id=str(span["trace_id"]),
+            event_id=event_id,
+            parent_span_id=str(span["parent_span_id"])
+            if span["parent_span_id"]
+            else "",
+            latency=f"{latency:.2f}",
+            latency_proportion="0",
+            attributes=span["attributes"],
+            childs=[],
+        )
+        store[str(span["span_id"])] = span_item
+
+    # normalize latency
+    for span_item in store.values():
+        span_item.latency_proportion = (
+            f"{float(span_item.latency) / max_latency * 100:.2f}"
+        )
+
+    # 2. build sequencial spans
+    root_items = []
+    for span_item in store.values():
+        if span_item.parent_span_id == "":
+            root_items.append(span_item)
+        else:
+            parent_span = store[span_item.parent_span_id]
+            parent_span.childs.append(span_item)
+
+    return root_items
+
+
+@app.get("/get_tracing")
+async def get_tracing() -> GetTracingResponse:
+    session_service = processor.short_term_memory.session_service
+    session = await session_service.get_session(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
+    )
+    if not session:
+        return GetTracingResponse(root_spans=[])
+    # ====== prevent existing file ======
+    tracing_file_path = processor.tracer.dump(
+        processor.runner.user_id, processor.session_id
+    )
+    os.remove(tracing_file_path)
+    # ====== ====================== ======
 
+    tracing_file_path = processor.tracer.dump(
+        processor.runner.user_id, processor.session_id
+    )
+    root_spans = parse_tracing_file(filepath=tracing_file_path)
+    os.remove(tracing_file_path)
+    return GetTracingResponse(root_spans=root_spans)
+
+
+@app.get("/get_event")
+async def get_event(event_id: str) -> GetEventResponse:
+    session_service = processor.short_term_memory.session_service
+    session = await session_service.get_session(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
+    )
 
-@app.post("/replace_prompt")
-async def replace_prompt(request: ReplacePromptRequest):
-    agent_processor.agent.instruction = request.prompt
-    return {"message": "Prompt replaced"}
+    # prevent no session created
+    if session:
+        for event in session.events:
+            if event.id == event_id:
+                return GetEventResponse(
+                    event=event.model_dump_json(exclude_none=True, by_alias=True)
+                )
+        return GetEventResponse(event=json.dumps({"info": "not an event span"}))
 
+    return GetEventResponse(event=json.dumps({"info": "not an event span"}))
 
-@app.get("/save_session")
-async def save_session(session_id: str):
-    """Save session to long term memory"""
-    await agent_processor.save_session_to_long_term_memory(session_id)
-    return {"message": "Session saved"}
 
+@app.get("/get_runner_config")
+async def get_runner_config() -> RunnerConfig:
+    return RunnerConfig(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
+    )
 
-@app.post("/update_prompt")
-async def update_prompt(request: ReplacePromptRequest):
-    agent_processor.agent.instruction = request.prompt
-    return {"message": "Prompt replaced"}
 
+@app.get("/set_runner_config")
+async def set_runner_config(
+    app_name: str, user_id: str, session_id: str
+) -> RunnerConfig:
+    processor.runner.app_name = app_name
+    processor.runner.user_id = user_id
+    processor.session_id = session_id
+    return RunnerConfig(
+        app_name=processor.runner.app_name,
+        user_id=processor.runner.user_id,
+        session_id=processor.session_id,
+    )
 
-@app.get("/history_sessions")
-async def get_history_sessions(session_id: str) -> GetHistorySessionsResponse:
-    events = await agent_processor.get_history_sessions(session_id)
-    return GetHistorySessionsResponse(events=events)
 
+@app.get("/set_prompt")
+async def set_prompt(prompt: str):
+    processor.runner.agent.instruction = prompt
+    return {"result": "success"}
 
-# run agent with user ID
-@app.post("/run")
-async def run(request: RunAgentRequest) -> RunAgentResponse:
-    agent_processor.agent.model._additional_args["stream"] = False
-    agent_processor.agent.model._additional_args["stream_options"] = {}
 
-    session_id = request.session_id
-    message = request.message
-    response = await agent_processor.runner.run_with_final_event(
-        messages=message,
-        session_id=session_id,
+@app.get("/refine_prompt")
+async def refine_prompt(feedback: str):
+    refined_prompt = ""  # prevent error processing
+    refined_prompt = processor.agent_pilot.optimize(
+        agents=[processor.agent], feedback=feedback
     )
-    return RunAgentResponse(event=response)
+    return {"prompt": refined_prompt}
 
 
-@app.get("/testcases")
-async def get_testcases(session_id: str):
-    testcases = await agent_processor.get_testcases(session_id)
-    return {"data": testcases}
+@app.get("/get_testcases")
+async def get_testcases():
+    test_cases = await processor.get_testcases()
+    return {"test_cases": test_cases}
 
 
-@app.post("/evaluate")
+@app.get("/evaluate")
 async def evaluate():
-    agent_processor.agent.model._additional_args["stream"] = False
-    agent_processor.agent.model._additional_args["stream_options"] = {}
-
-    test_cases = await agent_processor.evaluate()
-    return {"data": test_cases}
+    test_cases = await processor.evaluate()
+    return {"test_cases": test_cases}
 
 
-@app.get("/run_sse")
-async def run_sse(session_id: str, message: str):
-    return StreamingResponse(
-        agent_processor.run_sse(session_id, message),
-        media_type="text/event-stream",
-    )
-
-
-@app.get("/trace")
-def trace(session_id: str) -> TraceAgentResponse:
-    content = agent_processor.trace(
-        session_id=session_id,
-    )
-    return TraceAgentResponse(content=content)
-
-
-@app.get("/get/event")
-async def get_event(session_id: str, invocation_id: str) -> GetEventResponse:
-    event_str = await agent_processor.get_event(session_id, invocation_id)
-    return GetEventResponse(event=event_str)
+@app.get("/save_session")
+async def save_session():
+    await processor.runner.save_session_to_long_term_memory(processor.session_id)
+    return {"result": "success"}
 
 
 app.mount(
@@ -227,6 +380,13 @@ app.mount(
 )
 
 
-def get_fast_api_app(agents_dir: str) -> FastAPI:
-    agent_processor.agents_dir = agents_dir
+def get_fast_api_app(agent, short_term_memory):
+    global processor
+    processor = StudioProcessor(
+        app_name="veadk_studio",
+        user_id="studio_user",
+        session_id="studio_session",
+        agent=agent,
+        short_term_memory=short_term_memory,
+    )
     return app

veadk/cli/studio/models.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import BaseModel
+
+
+class GetAgentResponse(BaseModel):
+    name: str
+    description: str
+    instruction: str
+    model_name: str
+    short_term_memory_backend: str
+    short_term_memory_db_url: str
+    long_term_memory_backend: str
+    knowledgebase_backend: str
+
+
+class SpanItem(BaseModel):
+    name: str
+    span_id: str
+    trace_id: str
+    event_id: str
+    parent_span_id: str
+    latency: str
+    latency_proportion: str
+    attributes: dict
+    childs: list["SpanItem"]
+
+
+class GetTracingResponse(BaseModel):
+    root_spans: list[SpanItem]
+
+
+class GetEventResponse(BaseModel):
+    event: str
+
+
+class RunnerConfig(BaseModel):
+    app_name: str
+    user_id: str
+    session_id: str
+
+
+class TextPart(BaseModel):
+    type: str = "text"
+    state: str = "done"
+    text: str
+
+
+class ToolPart(BaseModel):
+    type: str
+    toolCallId: str
+    state: str = ""
+    input: dict = {}
+    output: dict = {}
+
+
+class Message(BaseModel):
+    id: str
+    role: str
+    parts: list[TextPart | ToolPart]
+
+
+class GetHistoryMessagesResponse(BaseModel):
+    messages: list[Message]