feat(demohouse/shopping): mock vdb

Author: liyuxuan-bd

demohouse/shopping/backend/code/main.py

@@ -1,16 +1,23 @@
 import json
 import logging
 import os
+import time
 from typing import AsyncIterable
 from arkitect.launcher.local.serve import launch_serve
 from arkitect.telemetry.trace import task
-from arkitect.types.llm.model import ArkChatRequest, ArkChatParameters
+from arkitect.types.llm.model import (
+    ArkChatRequest,
+    ArkChatParameters,
+    ArkChatCompletionChunk,
+    BotUsage,
+    ActionDetail,
+    ToolDetail,
+)
 from volcenginesdkarkruntime.types.chat import ChatCompletionChunk
 
 from arkitect.core.component.context.context import Context
 
 from arkitect.core.runtime import Response
-from volcenginesdkarkruntime.types.chat.chat_completion_chunk import Choice, ChoiceDelta
 
 from arkitect.core.component.context.model import State
 
@@ -19,13 +26,37 @@
 
 logger = logging.getLogger(__name__)
 
-DOUBAO_VLM_ENDPOINT = "doubao-1-5-vision-pro-32k-250115"
+DOUBAO_VLM_ENDPOINT = "doubao-1-5-pro-32k-250115"
+
+
+@task()
+def get_vector_search_result_chunk(result: str) -> ArkChatCompletionChunk:
+    return ArkChatCompletionChunk(
+        id="",
+        created=int(time.time()),
+        model="",
+        object="chat.completion.chunk",
+        choices=[],
+        bot_usage=BotUsage(
+            action_details=[
+                ActionDetail(
+                    name="vector_search",
+                    count=1,
+                    tool_details=[
+                        ToolDetail(
+                            name="vector_search", input=None, output=json.loads(result)
+                        )
+                    ],
+                )
+            ]
+        ),
+    )
 
 
 @task()
 async def default_model_calling(
     request: ArkChatRequest,
-) -> AsyncIterable[ChatCompletionChunk]:
+) -> AsyncIterable[ArkChatCompletionChunk | ChatCompletionChunk]:
     parameters = ArkChatParameters(**request.__dict__)
     image_urls = [
         content.get("image_url", {}).get("url", "")
@@ -36,6 +67,12 @@ async def default_model_calling(
     ]
     image_url = image_urls[-1] if len(image_urls) > 0 else ""
 
+    # only search for relevant product
+    if request.metadata and request.metadata.get("search"):
+        result = await vector_search("", image_url)
+        yield get_vector_search_result_chunk(result)
+        return
+
     async def modify_url_hook(
         state: State, param: ChatCompletionMessageToolCallParam
     ) -> ChatCompletionMessageToolCallParam:
@@ -55,12 +92,8 @@ async def modify_url_hook(
         async for chunk in stream:
             if tool_call and chunk.choices:
                 tool_result = ctx.get_latest_message()
-                chunk.choices.append(
-                    Choice(
-                        role="tool",
-                        delta=ChoiceDelta(content=tool_result.get("content")),
-                        index=len(chunk.choices),
-                    )
+                yield get_vector_search_result_chunk(
+                    str(tool_result.get("content", ""))
                 )
                 tool_call = False
             yield chunk