veadk-python 0.1.0__py3-none-any.whl

veadk/cli/studio/agent_processor.py

@@ -0,0 +1,247 @@
+# 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.
+
+
+import json
+import os
+
+from deepeval.metrics import GEval
+from deepeval.test_case import LLMTestCaseParams
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.cli.utils.agent_loader import AgentLoader
+from google.adk.models.llm_response import LlmResponse
+
+from veadk.agent import Agent
+from veadk.cli.services.agentpilot.agentpilot import AgentPilot
+from veadk.config import getenv
+from veadk.evaluation.deepeval_evaluator.deepeval_evaluator import DeepevalEvaluator
+from veadk.evaluation.eval_set_recorder import EvalSetRecorder
+from veadk.memory.short_term_memory import ShortTermMemory
+from veadk.runner import Runner
+from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter
+from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
+
+
+class AgentProcessor:
+    def __init__(self, agents_dir: str = ""):
+        self.agents_dir = agents_dir
+
+        self.agent: Agent = None
+
+        self.runner: Runner = None
+
+        self.prompt_optimizer: AgentPilot = AgentPilot(
+            api_key=getenv("AGENT_PILOT_API_KEY")
+        )
+
+        self.trace_dict = {}
+
+        self.short_term_memory = None
+
+        self.evaluator = None
+        self.evalset_recorder = None
+
+    def get_agent_names(self) -> list[str]:
+        agent_dirs = []
+
+        for dirpath, dirnames, filenames in os.walk(self.agents_dir):
+            required_files = {"agent.py", "__init__.py"}
+            existing_files = set(filenames)
+            if required_files.issubset(existing_files):
+                agent_dirs.append(dirpath)
+
+        return [agent_dir.split("/")[-1] for agent_dir in agent_dirs]
+
+    def _set_short_term_memory(self, backend: str, db_url: str):
+        return ShortTermMemory(backend=backend, db_url=db_url)
+
+    async def set_runner(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        short_term_memory_backend: str,
+        short_term_memory_db_url: str,
+    ):
+        short_term_memory = self._set_short_term_memory(
+            backend=short_term_memory_backend, db_url=short_term_memory_db_url
+        )
+        await short_term_memory.create_session(
+            app_name=app_name, user_id=user_id, session_id=session_id
+        )
+
+        self.short_term_memory = short_term_memory
+
+    def set_agent(
+        self,
+        agent_name: str,
+    ):
+        try:
+            agent_loader = AgentLoader(agents_dir=self.agents_dir)
+            agent = agent_loader.load_agent(agent_name)
+            self.agent = agent
+            self.agent.after_model_callback = after_model_callback
+            if not self.agent.tracers:
+                exporter = InMemoryExporter()
+                self.agent.tracers = [OpentelemetryTracer(exporters=[exporter])]
+            self.runner = Runner(
+                agent=self.agent, short_term_memory=self.short_term_memory
+            )
+            self.evalset_recorder = EvalSetRecorder(
+                self.short_term_memory.session_service
+            )
+        except Exception as e:
+            print(e)
+
+    def get_memory_status(self) -> dict:
+        short_term_memory_backend = self.runner.short_term_memory.backend
+        long_term_memory_backend = "No long term memory provided."
+        if self.agent.long_term_memory:
+            long_term_memory_backend = self.agent.long_term_memory.backend
+        return {
+            "short_term_memory_backend": short_term_memory_backend,
+            "long_term_memory_backend": long_term_memory_backend,
+        }
+
+    async def get_testcases(self, session_id: str):
+        try:
+            dump_path = await self.evalset_recorder.dump(
+                self.runner.app_name, self.runner.user_id, session_id
+            )
+            os.remove(dump_path)
+        except Exception as _:
+            pass
+
+        dump_path = await self.evalset_recorder.dump(
+            self.runner.app_name, self.runner.user_id, session_id
+        )
+        self._eval_set_path = dump_path
+
+        # 2. build
+        self.evaluator = DeepevalEvaluator(self.agent)
+        self.evaluator.generate_eval_data(dump_path)
+
+        test_cases = self.evaluator.get_data()
+
+        return test_cases
+
+    async def evaluate(self):
+        metrics = [
+            GEval(
+                name="Correctness&MatchDegree",
+                criteria="Judge the correctness and match degree of the model output with the expected output.",
+                evaluation_params=[
+                    LLMTestCaseParams.INPUT,
+                    LLMTestCaseParams.ACTUAL_OUTPUT,
+                    LLMTestCaseParams.EXPECTED_OUTPUT,
+                ],
+            ),
+        ]
+
+        await self.evaluator.eval(
+            eval_set_file_path=self._eval_set_path, metrics=metrics
+        )
+        test_cases = self.evaluator.get_data()
+
+        os.remove(self._eval_set_path)
+        self._eval_set_path = ""
+
+        return test_cases
+
+    def optimize_prompt(self, prompt: str, feedback: str) -> str:
+        self.agent.instruction = prompt
+        return self.prompt_optimizer.optimize(agents=[self.agent], feedback=feedback)
+
+    async def get_history_sessions(self, session_id: str) -> list[str]:
+        events = []
+
+        session_service = self.runner.short_term_memory.session_service
+        session = await session_service.get_session(
+            app_name=self.runner.app_name,
+            user_id=self.runner.user_id,
+            session_id=session_id,
+        )
+
+        # prevent no session created
+        if session:
+            for event in session.events:
+                event_string = event.model_dump_json(exclude_none=True, by_alias=True)
+                events.append(event_string)
+                print(event_string)
+            return events
+        else:
+            return []
+
+    async def get_event(self, session_id: str, invocation_id: str) -> str:
+        session_service = self.runner.short_term_memory.session_service
+        session = await session_service.get_session(
+            app_name=self.runner.app_name,
+            user_id=self.runner.user_id,
+            session_id=session_id,
+        )
+
+        # prevent no session created
+        if session:
+            for event in session.events:
+                if event.invocation_id == invocation_id:
+                    # find model response
+                    if (
+                        event.author != "user"
+                        and event.content.parts[0]
+                        and event.content.parts[0].text
+                    ):
+                        return event.model_dump_json(exclude_none=True, by_alias=True)
+            return ""
+        return ""
+
+    async def run_sse(self, session_id: str, prompt: str):
+        self.agent.model._additional_args["stream"] = True
+        self.agent.model._additional_args["stream_options"] = {"include_usage": True}
+        self.agent.after_model_callback = after_model_callback
+
+        if self.agent and self.runner:
+            async for chunk in self.runner.run_sse(
+                session_id=session_id, prompt=prompt
+            ):
+                yield chunk
+
+    def trace(self, session_id: str) -> str:
+        tracing_file_path = self.runner.save_tracing_file(session_id=session_id)
+
+        content = "No content!"
+        with open(tracing_file_path, "r", encoding="utf-8") as file:
+            data = json.load(file)
+            content = json.dumps(data, ensure_ascii=False, indent=2)
+
+        # remove file
+        if os.path.exists(tracing_file_path):
+            os.remove(tracing_file_path)
+
+        return content
+
+    async def save_session_to_long_term_memory(self, session_id: str):
+        """Save session to long term memory"""
+        await self.runner.save_session_to_long_term_memory(session_id)
+
+
+def after_model_callback(callback_context: CallbackContext, llm_response: LlmResponse):
+    if llm_response.content and llm_response.content.parts:
+        valid_parts = [
+            part
+            for part in llm_response.content.parts
+            if not (part.function_call and not part.function_call.name)
+        ]
+        if len(valid_parts) < len(llm_response.content.parts):
+            llm_response.content.parts = valid_parts
+    return llm_response

veadk/cli/studio/fast_api.py

@@ -0,0 +1,232 @@
+# 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.
+import os
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+
+from veadk.cli.studio.agent_processor import AgentProcessor
+from veadk.cli.studio.model import (
+    GetAgentResponse,
+    GetAgentsResponse,
+    GetEventResponse,
+    GetHistorySessionsResponse,
+    GetMemoryResponse,
+    OptimizePromptRequest,
+    OptimizePromptResponse,
+    ReplacePromptRequest,
+    RunAgentRequest,
+    RunAgentResponse,
+    SetAgentRequest,
+    SetAgentResponse,
+    SetRunnerRequest,
+    TraceAgentResponse,
+)
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    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")
+NEXT_HTML_DIR = NEXT_STATIC_DIR
+
+
+@app.get("/")
+async def read_root():
+    index_path = os.path.join(NEXT_HTML_DIR, "index.html")
+    if os.path.exists(index_path):
+        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,
+    )
+    # 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,
+    )
+
+
+# get current agent
+@app.get("/agent")
+async 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,
+    )
+
+
+@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)
+
+
+@app.post("/replace_prompt")
+async def replace_prompt(request: ReplacePromptRequest):
+    agent_processor.agent.instruction = request.prompt
+    return {"message": "Prompt replaced"}
+
+
+@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.post("/update_prompt")
+async def update_prompt(request: ReplacePromptRequest):
+    agent_processor.agent.instruction = request.prompt
+    return {"message": "Prompt replaced"}
+
+
+@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)
+
+
+# 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,
+    )
+    return RunAgentResponse(event=response)
+
+
+@app.get("/testcases")
+async def get_testcases(session_id: str):
+    testcases = await agent_processor.get_testcases(session_id)
+    return {"data": testcases}
+
+
+@app.post("/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}
+
+
+@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.mount(
+    "/_next",
+    StaticFiles(directory=os.path.join(NEXT_STATIC_DIR, "_next")),
+    name="next_static",
+)
+
+app.mount(
+    "/",
+    StaticFiles(directory=NEXT_STATIC_DIR, html=True),
+    name="static",
+)
+
+
+def get_fast_api_app(agents_dir: str) -> FastAPI:
+    agent_processor.agents_dir = agents_dir
+    return app

veadk/cli/studio/model.py

@@ -0,0 +1,116 @@
+# 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 GetAgentsResponse(BaseModel):
+    agents_dir: str
+    agents: list[str]
+
+
+class SetAgentRequest(BaseModel):
+    agent_name: str
+
+
+class SetAgentResponse(BaseModel):
+    name: str
+    description: str
+    model_name: str
+    instruction: str
+    long_term_memory_backend: str = ""
+    knowledgebase_backend: str = ""
+
+
+class SetRunnerRequest(BaseModel):
+    app_name: str
+    user_id: str
+    session_id: str
+    short_term_memory_backend: str
+    short_term_memory_db_url: str
+
+
+class GetMemoryResponse(BaseModel):
+    short_term_memory_backend: str
+    # short_term_memory_session_number: str
+    long_term_memory_backend: str
+    # long_term_memory_session_number: str
+
+
+class GetAgentResponse(BaseModel):
+    name: str
+    description: str
+    model_name: str
+    instruction: str
+
+
+class GetHistorySessionsResponse(BaseModel):
+    events: list[str]
+
+
+class OptimizePromptRequest(BaseModel):
+    prompt: str
+    feedback: str
+
+
+class OptimizePromptResponse(BaseModel):
+    prompt: str
+
+
+class ReplacePromptRequest(BaseModel):
+    prompt: str
+
+
+class RunAgentRequest(BaseModel):
+    session_id: str
+    message: str
+
+
+class RunAgentResponse(BaseModel):
+    event: str
+
+
+class TraceAgentResponse(BaseModel):
+    content: str
+
+
+class EvaluateAgentResponse(BaseModel):
+    content: str
+
+
+class GetEventResponse(BaseModel):
+    event: str
+
+
+# ========== messages ==========
+# class BotMessage(BaseModel):
+#     type: str = "bot"
+#     event_id: str
+#     invocation_id: str
+#     timestamp: any
+#     content: str
+#     prompt_tokens: str
+#     candidate_tokens: str
+
+
+# class ToolMessage(BaseModel):
+#     type: str = "tool"
+#     event_id: str
+#     invocation_id: str
+#     function_call_id: str
+#     tool_name: str
+#     tool_args: dict
+#     tool_response: str
+#     timestamp: any
+#     finish: bool

veadk/cloud/__init__.py

@@ -0,0 +1,13 @@
+# 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.