feat(demohouse/shopping): mock vdb

Author: liyuxuan-bd

demohouse/shopping/backend/code/main.py

@@ -0,0 +1,86 @@
+import json
+import logging
+import os
+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 volcenginesdkarkruntime.types.chat import (
+    ChatCompletionChunk,
+    ChatCompletionContentPartImageParam,
+)
+
+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
+
+from arkitect.types.llm.model import ChatCompletionMessageToolCallParam
+from vdb import vector_search
+
+logger = logging.getLogger(__name__)
+
+DOUBAO_VLM_ENDPOINT = "doubao-1-5-pro-32k-250115"
+
+
+@task()
+async def default_model_calling(
+    request: ArkChatRequest,
+) -> AsyncIterable[ChatCompletionChunk]:
+    parameters = ArkChatParameters(**request.__dict__)
+    image_urls = [
+        content.get("image_url", {}).get("url", "")
+        for message in request.messages
+        if isinstance(message.content, list)
+        for content in message.content
+        if isinstance(content, ChatCompletionContentPartImageParam)
+    ]
+    image_url = image_urls[-1] if len(image_urls) > 0 else ""
+
+    async def modify_url_hook(
+        state: State, param: ChatCompletionMessageToolCallParam
+    ) -> ChatCompletionMessageToolCallParam:
+        arguments = json.loads(param["function"]["arguments"])
+        arguments["image_url"] = image_url
+        param["function"]["arguments"] = json.dumps(arguments)
+        return param
+
+    async with Context(
+        model=DOUBAO_VLM_ENDPOINT, tools=[vector_search], parameters=parameters
+    ) as ctx:
+        ctx.tool_hooks.update(vector_search=[modify_url_hook])
+        stream = await ctx.completions.create(
+            messages=[m.model_dump() for m in request.messages], stream=True
+        )
+        tool_call = False
+        async for chunk in stream:
+            if tool_call:
+                tool_result = ctx.get_latest_message()
+                chunk.choices.append(
+                    Choice(
+                        role="tool",
+                        delta=ChoiceDelta(content=tool_result.get("content")),
+                        index=len(chunk.choices),
+                    )
+                )
+            yield chunk
+            if chunk.choices and chunk.choices[0].finish_reason == "tool_calls":
+                tool_call = True
+
+
+@task()
+async def main(request: ArkChatRequest) -> AsyncIterable[Response]:
+    async for resp in default_model_calling(request):
+        yield resp
+
+
+if __name__ == "__main__":
+    port = os.getenv("_FAAS_RUNTIME_PORT")
+    launch_serve(
+        package_path="main",
+        port=int(port) if port else 8888,
+        health_check_path="/v1/ping",
+        endpoint_path="/api/v3/bots/chat/completions",
+    )

demohouse/shopping/backend/code/vdb.py

@@ -0,0 +1,86 @@
+import os
+
+import tos
+
+from httpx import Timeout
+from tos import HttpMethodType
+from volcenginesdkarkruntime import AsyncArk
+from volcenginesdkarkruntime.types.multimodal_embedding import (
+    MultimodalEmbeddingContentPartTextParam,
+    MultimodalEmbeddingResponse,
+    MultimodalEmbeddingContentPartImageParam,
+)
+from volcengine.viking_db import *
+from volcenginesdkarkruntime.types.multimodal_embedding.embedding_content_part_image_param import (
+    ImageURL,
+)
+
+
+COLLECTION_NAME = "shopping_demo"
+INDEX_NAME = "shopping_demo"
+MODEL_NAME = "doubao-embedding-vision-241215"
+LIMIT = 6
+SCORE_THRESHOLD = 300
+
+vikingdb_service = VikingDBService(
+    host="api-vikingdb.volces.com",
+    region="cn-beijing",
+    scheme="https",
+    connection_timeout=30,
+    socket_timeout=30,
+)
+vikingdb_service.set_ak(os.environ.get("VOLC_ACCESSKEY"))
+vikingdb_service.set_sk(os.environ.get("VOLC_SECRETKEY"))
+
+tos_client = tos.TosClientV2(
+    os.getenv("VOLC_ACCESSKEY"),
+    os.getenv("VOLC_SECRETKEY"),
+    "tos-cn-beijing.volces.com",
+    "cn-beijing",
+)
+
+
+async def vector_search(text: str, image_url: str) -> str:
+    """获取商品相关信息，当想要了解商品信息，比如价格，详细介绍，销量，评价时调用该工具
+
+    Args:
+        text: 商品的描述信息
+        image_url: 固定填写为<image_url>
+    """
+    client = AsyncArk(timeout=Timeout(connect=1.0, timeout=60.0))
+    embedding_input = [MultimodalEmbeddingContentPartTextParam(type="text", text=text)]
+    if image_url != "":
+        embedding_input.append(
+            MultimodalEmbeddingContentPartImageParam(
+                type="image_url", image_url=ImageURL(url=image_url)
+            )
+        )
+    resp: MultimodalEmbeddingResponse = await client.multimodal_embeddings.create(
+        model=MODEL_NAME,
+        input=embedding_input,
+    )
+    embedding = resp.data.get("embedding", [])
+    index = await vikingdb_service.async_get_index(COLLECTION_NAME, INDEX_NAME)
+    retrieve = await index.async_search_by_vector(vector=embedding, limit=LIMIT)
+    retrieve_fields = [
+        json.loads(result.fields.get("data"))
+        for result in retrieve
+        if result.score > SCORE_THRESHOLD
+    ]
+    mock_data = [
+        {
+            "名称": item.get("Name", ""),
+            "类别": item.get("category", ""),
+            "子类别": item.get("sub_category", ""),
+            "价格": item.get("price", "99"),
+            "销量": item.get("sales", "999"),
+            "商品链接": tos_client.pre_signed_url(
+                http_method=HttpMethodType.Http_Method_Get,
+                bucket="shopping",
+                key=item.get("key", ""),
+                expires=600,
+            ).signed_url,
+        }
+        for item in retrieve_fields
+    ]
+    return json.dumps(mock_data, ensure_ascii=False)