Merge pull request #9 from dingzhaohan/main

fix job add, tiefblue upload interface

Author: dingzhaohan

pkg/sumdb/sum.golang.org/latest

@@ -0,0 +1,5 @@
+go.sum database tree
+39073249
+CHoGLc3xiTLWJIkUvqwtoUC8uUi8tsi/AMKPaA1nDxo=
+
+— sum.golang.org Az3grieYv97xelotCMbg3+3zXebCIZkZfor40fTCTWo4rNjv+mQ8eEJFgiyyjPkL1h6GWziyoBzovPY6pFcfbssUNQU=

setup.py

@@ -10,8 +10,8 @@ def read_file(filename):
 
 
 setup(
-    name="bohrium",
-    version="0.1.0",
+    name="bohrium-sdk",
+    version="0.13.0",
     author="dingzhaohan",
     author_email="[email protected]",
     url="https://github.com/dingzhaohan",
@@ -24,7 +24,15 @@ def read_file(filename):
         # include .txt all of them
         "": ["*.txt"]
     },
-    install_requires=[],
+    install_requires=[
+        "distro",
+        "httpx",
+        "typing_extensions",
+        "anyio",
+        "pyhumps",
+        "requests",
+        "tqdm",
+    ],
     python_requires=">=3.7",
     entry_points={},
 )

src/bohrium/_base_client.py

@@ -143,6 +143,8 @@ def _build_headers(self, custom_headers) -> httpx.Headers:
         headers_dict = _merge_mappings(
             self.default_headers, self._custom_headers, custom_headers
         )
+        # 过滤掉 value 为 None 的 header
+        headers_dict = {k: v for k, v in headers_dict.items() if v is not None}
         headers = httpx.Headers(headers_dict)
         return headers or dict()
 
@@ -177,17 +179,45 @@ def platform_headers(self) -> Dict[str, str]:
         exceptions=(httpx.RequestError,),
     )
     def _request(
-        self, method: str, path: str, json=None, headers=None, **kwargs
+        self, method: str, path: str, json=None, headers=None, data=None, **kwargs
     ) -> httpx.Response:
         url = urljoin(str(self._base_url), path)
         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(),
                 url,
                 json=json,
+                data=data,
                 headers=merged_headers,
                 params=merged_params,
             )

src/bohrium/_client.py

@@ -14,14 +14,19 @@
 
 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,
         access_key: Optional[str] = None,
+        app_key: Optional[str] = None,
         base_url: Optional[Union[str, URL]] = None,
         project_id: Optional[str] = None,
         timeout: Optional[Union[float, Timeout]] = 30.0,
@@ -36,15 +41,11 @@ def __init__(
                 "The api_key client option must be set either by passing api_key to the client or by setting the ACCESS_KEY environment variable"
             )
         self.access_key = access_key
-
+        self.app_key = app_key or os.environ.get("BOHRIUM_APP_KEY")
+        self.params = {"accessKey": self.access_key}
         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:
@@ -63,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
@@ -71,6 +76,7 @@ def default_headers(self) -> dict[str, str]:
             "Accept": "application/json",
             "Content-Type": "application/json",
             "Authorization": f"Bearer {self.access_key}",
+            "x-app-key": self.app_key,
         }
 
     def _make_status_error(

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"
+    "SigmaSearch", "AsyncSigmaSearch", 
+    "UniParser", "AsyncUniParser",
+    "KnowledgeBase", "AsyncKnowledgeBase",
+    "Paper", "AsyncPaper"
+]
 
-__all__ = ["Job", "AsyncJob"]

src/bohrium/resources/job/job.py

@@ -2,7 +2,7 @@
 import os
 import uuid
 from pathlib import Path
-
+import humps
 # from ..._resource import BaseClient
 from pprint import pprint
 from typing import Optional
@@ -18,6 +18,21 @@
 
 class Job(SyncAPIResource):
 
+    def create(self, project_id, name='', group_id=0):
+        data = {
+            'projectId': project_id
+        }
+        if name:
+            data['name'] = name
+        if group_id:
+            data['bohrGroupId'] = group_id
+        try:
+            data = self._client.post(f'/openapi/v1/job/create', json=data, params=self._client.params)
+            data = data.json()
+        except Exception as e:
+            raise e
+        return data.get("data", {})
+    
     def detail(self, job_id):
         log.info(f"detail job {job_id}")
         response = self._client.get(f"/openapi/v1/job/{job_id}")
@@ -43,7 +58,6 @@ def submit(
     ):
         # log.info(f"submit job {name},project_id:{project_id}")
         data = self.create_job(project_id, job_name, job_group_id)
-        print(data)
         if work_dir != "":
             if not os.path.exists(work_dir):
                 raise FileNotFoundError
@@ -73,38 +87,39 @@ def submit(
         )
         return self.insert(job_add_request.to_dict())
 
-    def insert(self, data):
-        # log.info(f"insert job {data}")
-        response = self._client.post("/openapi/v2/job/add", json=data)
-        pprint(response.request)
-        print(response.json())
+    def insert(self, **kwargs):
+        camel_data = {humps.camelize(k): v for k, v in kwargs.items()}
+        if not isinstance(camel_data['ossPath'], list):
+            camel_data['ossPath'] = [camel_data['ossPath']]
+        if 'logFile' in camel_data:
+            camel_data['logFiles'] = camel_data['logFile']
+        if 'logFiles' in camel_data and not isinstance(camel_data['logFiles'], list):
+            camel_data['logFiles'] = [camel_data['logFiles']]
+        response = self._client.post("/openapi/v2/job/add", json=camel_data)
+        return response.json().get("data")
 
     def delete(self, job_id):
         # log.info(f"delete job {job_id}")
         response = self._client.post(f"/openapi/v1/job/del/{job_id}")
-        pprint(response.request)
-        print(response.json())
+
 
     def terminate(self, job_id):
         # log.info(f"terminate job {job_id}")
         response = self._client.post(f"/openapi/v1/job/terminate/{job_id}")
-        pprint(response.request)
-        print(response.json())
+
 
     def kill(self, job_id):
         # log.info(f"kill job {job_id}")
         response = self._client.post(f"/openapi/v1/job/kill/{job_id}")
-        pprint(response.request)
-        print(response.json())
+    
 
     def log(self, job_id, log_file="STDOUTERR", page=-1, page_size=8192):
         # log.info(f"log job {job_id}")
         response = self._client.get(
             f"/openapi/v1/job/{job_id}/log",
             params={"logFile": log_file, "page": page, "pageSize": page_size},
         )
-        pprint(response.request)
-        print(response.json().get("data")["log"])
+
         return response.json().get("data")["log"]
 
     def create_job(
@@ -132,8 +147,7 @@ def create_job(
             "bohrGroupId": group_id,
         }
         response = self._client.post(f"/openapi/v1/job/create", json=data)
-        pprint(response.request)
-        print(response.json())
+
         return response.json().get("data")
 
     def create_job_group(self, project_id, job_group_name):
@@ -142,9 +156,8 @@ def create_job_group(self, project_id, job_group_name):
             "/openapi/v1/job_group/add",
             json={"name": job_group_name, "projectId": project_id},
         )
-        pprint(response.request)
-        print(response.json())
-
+        return response.json().get("data")
+    
     def upload(
         self,
         file_path: str,

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