weknora

src/langbot/libs/weknora_api/v1/__init__.py

@@ -0,0 +1,4 @@
+from .client import AsyncWeKnoraClient
+from .errors import WeKnoraAPIError
+
+__all__ = ['AsyncWeKnoraClient', 'WeKnoraAPIError']

src/langbot/libs/weknora_api/v1/client.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import httpx
+import typing
+import json
+
+from .errors import WeKnoraAPIError
+
+
+class AsyncWeKnoraClient:
+    """WeKnora API 客户端"""
+
+    api_key: str
+    base_url: str
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = 'http://localhost:80/api/v1',
+    ) -> None:
+        self.api_key = api_key
+        self.base_url = base_url
+
+    async def create_session(
+        self,
+        title: str = '',
+        description: str = '',
+        timeout: float = 30.0,
+    ) -> str:
+        """创建会话，返回 session_id"""
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+            payload: dict[str, typing.Any] = {}
+            if title:
+                payload['title'] = title
+            if description:
+                payload['description'] = description
+
+            response = await client.post(
+                '/sessions',
+                headers={
+                    'X-API-Key': self.api_key,
+                    'Content-Type': 'application/json',
+                },
+                json=payload,
+            )
+
+            if response.status_code not in (200, 201):
+                raise WeKnoraAPIError(f'{response.status_code} {response.text}')
+
+            data = response.json()
+            return data['data']['id']
+
+    async def agent_chat(
+        self,
+        session_id: str,
+        query: str,
+        user: str,
+        agent_id: str = '',
+        knowledge_base_ids: list[str] | None = None,
+        web_search_enabled: bool = False,
+        timeout: float = 120.0,
+    ) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
+        """
+        Agent 智能对话（SSE 流式）
+
+        响应事件类型:
+        - agent_query: Agent 开始处理
+        - thinking: 思考过程
+        - tool_call: 工具调用
+        - tool_result: 工具结果
+        - references: 知识库引用
+        - answer: 回答内容
+        - reflection: 反思
+        - session_title: 会话标题
+        - error: 错误
+        """
+        if knowledge_base_ids is None:
+            knowledge_base_ids = []
+
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+            payload: dict[str, typing.Any] = {
+                'query': query,
+                'agent_enabled': True,
+                'channel': 'im',
+            }
+            if agent_id:
+                payload['agent_id'] = agent_id
+            if knowledge_base_ids:
+                payload['knowledge_base_ids'] = knowledge_base_ids
+            if web_search_enabled:
+                payload['web_search_enabled'] = True
+
+            async with client.stream(
+                'POST',
+                f'/agent-chat/{session_id}',
+                headers={
+                    'X-API-Key': self.api_key,
+                    'Content-Type': 'application/json',
+                },
+                json=payload,
+            ) as r:
+                async for chunk in r.aiter_lines():
+                    if r.status_code != 200:
+                        raise WeKnoraAPIError(f'{r.status_code} {chunk}')
+                    if chunk.strip() == '':
+                        continue
+                    if chunk.startswith('data:'):
+                        try:
+                            yield json.loads(chunk[5:].strip())
+                        except json.JSONDecodeError:
+                            continue
+
+    async def knowledge_chat(
+        self,
+        session_id: str,
+        query: str,
+        user: str,
+        agent_id: str = 'builtin-quick-answer',
+        knowledge_base_ids: list[str] | None = None,
+        timeout: float = 120.0,
+    ) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
+        """
+        知识库 RAG 问答（SSE 流式）
+
+        响应事件类型:
+        - references: 知识库引用
+        - answer: 回答内容
+        """
+        if knowledge_base_ids is None:
+            knowledge_base_ids = []
+
+        async with httpx.AsyncClient(
+            base_url=self.base_url,
+            trust_env=True,
+            timeout=timeout,
+        ) as client:
+            payload: dict[str, typing.Any] = {
+                'query': query,
+                'channel': 'im',
+            }
+            if agent_id:
+                payload['agent_id'] = agent_id
+            if knowledge_base_ids:
+                payload['knowledge_base_ids'] = knowledge_base_ids
+
+            async with client.stream(
+                'POST',
+                f'/knowledge-chat/{session_id}',
+                headers={
+                    'X-API-Key': self.api_key,
+                    'Content-Type': 'application/json',
+                },
+                json=payload,
+            ) as r:
+                async for chunk in r.aiter_lines():
+                    if r.status_code != 200:
+                        raise WeKnoraAPIError(f'{r.status_code} {chunk}')
+                    if chunk.strip() == '':
+                        continue
+                    if chunk.startswith('data:'):
+                        try:
+                            yield json.loads(chunk[5:].strip())
+                        except json.JSONDecodeError:
+                            continue

src/langbot/libs/weknora_api/v1/errors.py

@@ -0,0 +1,6 @@
+class WeKnoraAPIError(Exception):
+    """WeKnora API 请求失败"""
+
+    def __init__(self, message: str = ''):
+        self.message = message
+        super().__init__(self.message)