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

demohouse/shopping/backend/code/vdb.py

@@ -48,7 +48,9 @@ async def vector_search(text: str, image_url: str) -> str:
         image_url: 固定填写为<image_url>
     """
     client = AsyncArk(timeout=Timeout(connect=1.0, timeout=60.0))
-    embedding_input = [MultimodalEmbeddingContentPartTextParam(type="text", text=text)]
+    embedding_input = []
+    if text != "":
+        embedding_input = [MultimodalEmbeddingContentPartTextParam(type="text", text=text)]
     if image_url != "":
         embedding_input.append(
             MultimodalEmbeddingContentPartImageParam(
@@ -74,7 +76,7 @@ async def vector_search(text: str, image_url: str) -> str:
             "子类别": item.get("sub_category", ""),
             "价格": item.get("price", "99"),
             "销量": item.get("sales", "999"),
-            "商品链接": tos_client.pre_signed_url(
+            "图片链接": tos_client.pre_signed_url(
                 http_method=HttpMethodType.Http_Method_Get,
                 bucket="shopping",
                 key=item.get("key", ""),