Merge branch 'main' into sandbox-env

Author: zjgemi

src/bohrium/_base_client.py

@@ -185,6 +185,33 @@ def _request(
         logger.info(f"Requesting {method} {url}")
         merged_headers = self._build_headers(headers)
         merged_params = self._build_params(kwargs.get("params"))
+        
+        # 处理文件上传
+        request_kwargs = {
+            "method": method.upper(),
+            "url": url,
+            "params": merged_params,
+        }
+        
+        # 处理超时参数
+        if "timeout" in kwargs:
+            request_kwargs["timeout"] = kwargs["timeout"]
+        
+        if json is not None:
+            request_kwargs["json"] = json
+            request_kwargs["headers"] = merged_headers
+        elif "files" in kwargs:
+            # 当有files参数时，不使用json参数，而是使用files和data
+            # 不设置headers，让httpx自动处理multipart/form-data
+            request_kwargs["files"] = kwargs["files"]
+            if "data" in kwargs:
+                request_kwargs["data"] = kwargs["data"]
+        elif "data" in kwargs:
+            request_kwargs["data"] = kwargs["data"]
+            request_kwargs["headers"] = merged_headers
+        else:
+            request_kwargs["headers"] = merged_headers
+            
         try:
             return self._client.request(
                 method.upper(),

src/bohrium/_client.py

@@ -14,10 +14,14 @@
 
 class Bohrium(SyncAPIClient):
     job: resources.Job
+    sigma_search: resources.SigmaSearch
+    uni_parser: resources.UniParser
+    knowledge_base: resources.KnowledgeBase
+    paper: resources.Paper
 
     # client options
     access_key: str
-    project_id: Union[str, None]
+    project_id: Optional[str]
 
     def __init__(
         self,
@@ -42,11 +46,6 @@ def __init__(
         if project_id is None:
             project_id = os.environ.get("BOHRIUM_PROJECT_ID")
 
-        if project_id is None:
-            raise BohriumError(
-                "The project_id client option must be set either by passing project_id to the client or by setting the BOHRIUM_PROJECT_ID environment variable"
-            )
-
         self.project_id = project_id
 
         if base_url is None:
@@ -65,6 +64,10 @@ def __init__(
         )
 
         self.job = resources.Job(self)
+        self.sigma_search = resources.SigmaSearch(self)
+        self.uni_parser = resources.UniParser(self)
+        self.knowledge_base = resources.KnowledgeBase(self)
+        self.paper = resources.Paper(self)
 
     @property
     @override

src/bohrium/resources/__init__.py

@@ -1,4 +1,16 @@
 
 from .job import Job, AsyncJob
+from .sigma_search import SigmaSearch, AsyncSigmaSearch
+from .uni_parser import UniParser, AsyncUniParser
+from .knowledge_base import KnowledgeBase, AsyncKnowledgeBase
+from .paper import Paper, AsyncPaper
 from .tiefblue import Tiefblue
-__all__ = ["Job", "AsyncJob", "Tiefblue"]
+
+__all__ = [
+    "Job", "AsyncJob", "Tiefblue"
+    "SigmaSearch", "AsyncSigmaSearch", 
+    "UniParser", "AsyncUniParser",
+    "KnowledgeBase", "AsyncKnowledgeBase",
+    "Paper", "AsyncPaper"
+]
+

src/bohrium/resources/knowledge_base/__init__.py

@@ -0,0 +1,3 @@
+from .knowledge_base import KnowledgeBase, AsyncKnowledgeBase
+
+__all__ = ["KnowledgeBase", "AsyncKnowledgeBase"]

src/bohrium/resources/knowledge_base/knowledge_base.py

@@ -0,0 +1,114 @@
+import logging
+from typing import Optional, List, Dict, Any, Union
+from pprint import pprint
+
+from ..._resource import AsyncAPIResource, SyncAPIResource
+from ..._response import APIResponse
+from ...types.knowledge_base.knowledge_base import (
+    HybridRecallRequest,
+    PaperRecallRequest,
+    PaperInfo,
+    ChunkSearchRequest
+)
+
+log = logging.getLogger(__name__)
+
+
+class KnowledgeBase(SyncAPIResource):
+    """知识库相关接口"""
+
+    def hybrid_recall(
+        self,
+        knowledge_base_id: int,
+        text: str,
+        k: int = 200,
+        keywords: Optional[Dict[str, float]] = None,
+        **kwargs
+    ):
+        """知识库混合召回"""
+        log.info(f"hybrid recall from knowledge base: {knowledge_base_id}")
+        
+        data = {
+            "knowledge_base_id": knowledge_base_id,
+            "text": text,
+            "k": k
+        }
+        
+        if keywords:
+            data["keywords"] = keywords
+        if kwargs:
+            data.update(kwargs)
+            
+        response = self._client.post("/openapi/v1/knowledge/recall/hybrid", json=data)
+        log.info(response.json())
+        return APIResponse(response).json.get("data")
+
+    def paper_recall(
+        self,
+        text: str,
+        k: int,
+        papers: List[Dict[str, str]],
+        **kwargs
+    ):
+        """单篇论文召回"""
+        log.info(f"paper recall: {len(papers)} papers")
+        
+        data = {
+            "text": text,
+            "k": k,
+            "papers": papers
+        }
+        
+        if kwargs:
+            data.update(kwargs)
+            
+        response = self._client.post("/openapi/v1/knowledge/recall/papers", json=data)
+        log.info(response.json())
+        return APIResponse(response).json.get("data")
+
+    def get_file_tree(
+        self,
+        folder_id: str,
+        **kwargs
+    ):
+        """获取单篇切片文件树"""
+        log.info(f"get file tree for folder: {folder_id}")
+        
+        params = {"folderId": folder_id}
+        if kwargs:
+            params.update(kwargs)
+            
+        response = self._client.get("/openapi/v1/knowledge/folder/file_tree", params=params)
+        log.info(response.json())
+        return APIResponse(response).json.get("data")
+
+    def search_by_md5_paper_id(
+        self,
+        md5: str,
+        paper_id: str = "",
+        page_num: int = 0,
+        page_size: int = 9999,
+        **kwargs
+    ):
+        """根据md5和paper_id搜索chunk信息"""
+        log.info(f"search chunk by md5: {md5}, paper_id: {paper_id}")
+        
+        data = {
+            "md5": md5,
+            "paper_id": paper_id,
+            "page_num": page_num,
+            "page_size": page_size
+        }
+        
+        if kwargs:
+            data.update(kwargs)
+            
+        response = self._client.post("/openapi/v1/knowledge/box/search_by_md5_paper_id", json=data)
+        log.info(response.json())
+        return APIResponse(response).json.get("data")
+
+
+
+class AsyncKnowledgeBase(AsyncAPIResource):
+    """异步知识库相关接口"""
+    pass

src/bohrium/resources/paper/__init__.py

@@ -0,0 +1,3 @@
+from .paper import Paper, AsyncPaper
+
+__all__ = ["Paper", "AsyncPaper"]

src/bohrium/resources/paper/paper.py

@@ -0,0 +1,52 @@
+import logging
+from typing import Optional, List, Dict, Any
+from pprint import pprint
+
+from ..._resource import AsyncAPIResource, SyncAPIResource
+from ..._response import APIResponse
+from ...types.paper.paper import PaperRAGRequest
+
+log = logging.getLogger(__name__)
+
+
+class Paper(SyncAPIResource):
+    """论文相关接口"""
+
+    def rag_pass_keyword(
+        self,
+        type: int,
+        rerank: int,
+        question: str,
+        page_size: int,
+        words: Optional[List[str]] = None,
+        start_time: Optional[str] = None,
+        end_time: Optional[str] = None,
+        **kwargs
+    ):
+        """论文RAG关键词检索"""
+        log.info(f"paper rag pass keyword: type={type}, rerank={rerank}")
+        
+        data = {
+            "type": type,
+            "rerank": rerank,
+            "question": question,
+            "pageSize": page_size
+        }
+        
+        if words:
+            data["words"] = words
+        if start_time:
+            data["startTime"] = start_time
+        if end_time:
+            data["endTime"] = end_time
+        if kwargs:
+            data.update(kwargs)
+            
+        response = self._client.post("/openapi/v1/paper/rag/pass/keyword", json=data)
+        log.info(response.json())
+        return APIResponse(response).json.get("data")
+
+
+class AsyncPaper(AsyncAPIResource):
+    """异步论文相关接口"""
+    pass

src/bohrium/resources/sigma_search/__init__.py

@@ -0,0 +1,3 @@
+from .sigma_search import SigmaSearch, AsyncSigmaSearch
+
+__all__ = ["SigmaSearch", "AsyncSigmaSearch"]