davidkhala.ai 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

davidkhala/ai/agent/dify/api/__init__.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Callable, Any, Optional
+from typing import Iterable, Callable, Any
 
 from davidkhala.utils.http_request import Request
 
@@ -14,7 +14,7 @@ class Iterator(Iterable):
     def __iter__(self):
         return self
 
-    def __init__(self, get_fn: Callable[[int, int], Any], r: Optional[dict]):
+    def __init__(self, get_fn: Callable[[int, int], Any], r: dict|None):
         self.response = r
         self.fn = get_fn
 

davidkhala/ai/agent/dify/api/app.py

@@ -77,12 +77,16 @@ class Conversation(API):
         r: Conversation.ChatResult = {
             'thought': [],
         }
-        for data in as_sse(response):
-            match data['event']:
-                case 'agent_thought':
-                    r['thought'].append(data['thought'])
-                case 'message_end':
-                    r['metadata'] = data['metadata']
+
+        for line in response.iter_lines():
+            if line and line!=b'event: ping':
+                data = json.loads(line[5:].decode())
+                match data['event']:
+                    case 'agent_thought':
+                        r['thought'].append(data['thought'])
+                    case 'message_end':
+                        r['metadata'] = data['metadata']
+
         return r
 
     def agent_chat(self, template: str, **kwargs) -> ChatResult:

davidkhala/ai/agent/dify/api/knowledge/chunk.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from davidkhala.ai.agent.dify.api import API
+from davidkhala.ai.agent.dify.api.knowledge.model import DocumentModel
+
+
+class Chunk(API):
+    def __init__(self, d: DocumentModel, segment_id: str):
+        super().__init__(d.api_key, f"{d.base_url}/segments/{segment_id}")
+
+    def get(self):
+        r = self.request(self.base_url, "GET")
+        assert r['doc_form']  # optional value text_model
+        return r['data']

davidkhala/ai/agent/dify/api/knowledge/dataset.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Iterable
+from urllib.parse import urlparse
+
+import requests
+
+from davidkhala.ai.agent.dify.api import API, Iterator
+from davidkhala.ai.agent.dify.api.knowledge.model import DatasetModel, DocumentModel
+
+
+class Dataset(API):
+    def __init__(self, api_key: str, base_url="https://api.dify.ai/v1"):
+        super().__init__(api_key, f"{base_url}/datasets")
+
+    def paginate_datasets(self, page=1, size=20):
+        r = self.request(self.base_url, "GET", params={
+            'page': page,
+            'limit': size,
+        })
+        return r
+
+    def list_datasets(self) -> Iterable[list[DatasetModel]]:
+        return Iterator(self.paginate_datasets, None)
+
+    @property
+    def ids(self):
+        for sub_list in self.list_datasets():
+            for dataset in sub_list:
+                yield dataset['id']
+
+    class Instance(API):
+        def __init__(self, d: Dataset, dataset_id: str):
+            super().__init__(d.api_key, f"{d.base_url}/{dataset_id}")
+
+        def get(self)-> DatasetModel:
+            d = self.request(self.base_url, "GET")
+            return DatasetModel.model_validate(d)
+
+        def upload(self, filename, *, path=None, url=None, document_id=None):
+            """
+            don't work for .html
+            work for .md
+            """
+            files = {}
+            if path:
+                with open(path, 'rb') as f:
+                    content = f.read()
+                if not filename:
+                    filename = os.path.basename(path)
+            elif url:
+                r = requests.get(url)
+                r.raise_for_status()
+                if not filename:
+                    parsed_url = urlparse(url)
+                    filename = Path(parsed_url.path).name
+                content = r.content
+            files['file'] = (filename, content)
+            if document_id:
+                # don't work for html
+                r = requests.post(f"{self.base_url}/documents/{document_id}/update-by-file", files=files,
+                                  **self.options)
+            else:
+                r = requests.post(f"{self.base_url}/document/create-by-file", files=files, **self.options)
+            r = self.on_response(r)
+            return r['document']
+
+        def paginate_documents(self, page=1, size=20):
+            return self.request(f"{self.base_url}/documents", "GET", params={
+                'page': page,
+                'limit': size
+            })
+
+        def list_documents(self) -> Iterable[DocumentModel]:
+            for document_batch in Iterator(self.paginate_documents, None):
+                for document in document_batch:
+                    yield DocumentModel(**document)
+
+        def has_document(self, name) -> bool:
+            return any(name == item['name'] for row in self.list_documents() for item in row)

davidkhala/ai/agent/dify/api/knowledge/document.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+import requests
+
+from davidkhala.ai.agent.dify.api import API, Iterator
+from davidkhala.ai.agent.dify.api.knowledge.dataset import Dataset
+from davidkhala.ai.agent.dify.api.knowledge.model import ChunkDict
+
+
+class Document(API):
+    def __init__(self, d: Dataset.Instance, document_id: str):
+        super().__init__(d.api_key, f"{d.base_url}/documents/{document_id}")
+
+    def exist(self):
+        try:
+            self.get()
+            return True
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 404:
+                return False
+            else:
+                raise e
+
+    def get(self):
+        return self.request(self.base_url, "GET")
+
+    def paginate_chunks(self, page=1, size=20):
+        return self.request(f"{self.base_url}/segments", "GET", params={
+            'page': page,
+            'limit': size
+        })
+
+    def list_chunks(self) -> Iterable[ChunkDict]:
+        for chunk_batch in Iterator(self.paginate_chunks, None):
+            for chunk in chunk_batch:
+                yield chunk
+
+    def delete(self):
+        if self.exist():
+            self.request(self.base_url, "DELETE")

davidkhala/ai/agent/dify/api/knowledge/model.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import List, Any, TypedDict
+
+from pydantic import BaseModel
+
+from davidkhala.ai.agent.dify.model.knowledge import Document as DocumentBase
+
+
+class RerankingModel(BaseModel):
+    reranking_provider_name: str | None
+    reranking_model_name: str | None
+
+
+class KeywordSetting(BaseModel):
+    keyword_weight: float
+
+
+class VectorSetting(BaseModel):
+    vector_weight: float
+    embedding_model_name: str
+    embedding_provider_name: str
+
+
+class Weights(BaseModel):
+    weight_type: str | None
+    keyword_setting: KeywordSetting | None
+    vector_setting: VectorSetting | None
+
+
+class RetrievalModelDict(BaseModel):
+    search_method: str
+    reranking_enable: bool
+    reranking_mode: str | None
+    reranking_model: RerankingModel | None
+    weights: Weights | None
+    top_k: int | None
+    score_threshold_enabled: bool
+    score_threshold: float
+
+
+class ExternalKnowledgeInfo(BaseModel):
+    external_knowledge_id: str | None
+    external_knowledge_api_id: str | None
+    external_knowledge_api_name: str | None
+    external_knowledge_api_endpoint: str | None
+
+
+class ExternalRetrievalModel(BaseModel):
+    top_k: int
+    score_threshold: float
+    score_threshold_enabled: bool
+
+
+class IconInfo(BaseModel):
+    icon_type: str
+    icon: str | None
+    icon_background: str | None
+    icon_url: str | None
+
+
+class DatasetModel(BaseModel):
+    id: str
+    name: str
+    description: str | None = None
+    provider: str
+    permission: str
+    data_source_type: str | None = None
+    indexing_technique: str | None = None
+    app_count: int
+    document_count: int
+    word_count: int
+    created_by: str | None = None
+    author_name: str | None = None
+    created_at: int | None = None
+    updated_by: str | None = None
+    updated_at: int | None = None
+    embedding_model: str | None = None
+    embedding_model_provider: str | None = None
+    embedding_available: bool
+    retrieval_model_dict: RetrievalModelDict | None = None
+    tags: List[Any] = []
+    doc_form: str | None = None
+    external_knowledge_info: ExternalKnowledgeInfo | None = None
+    external_retrieval_model: ExternalRetrievalModel | None = None
+    doc_metadata: List[Any] = []
+    built_in_field_enabled: bool
+    pipeline_id: str | None = None
+    runtime_mode: str | None = None
+    chunk_structure: str | None = None
+    icon_info: IconInfo | None = None
+    is_published: bool
+    total_documents: int | None = None
+    total_available_documents: int | None = None
+    enable_api: bool
+    is_multimodal: bool
+
+
+class DocumentModel(DocumentBase):
+    data_source_info: dict[str, str]
+    data_source_detail_dict: dict[str, dict]
+    dataset_process_rule_id: str
+    created_from: str
+    created_by: str
+    created_at: int
+    tokens: int
+    archived: bool
+    display_status: str
+    word_count: int
+    hit_count: int
+    doc_form: str
+    doc_metadata: dict
+    disabled_at: int
+    disabled_by: str
+
+
+class ChunkDict(TypedDict):
+    id: str
+    position: int
+    document_id: str
+    content: str
+    sign_content: str  # trimmed version of content
+    answer: str | None  # only used in QA chunk
+    word_count: int
+    tokens: int
+    keywords: list[str] | None
+    index_node_id: str  # chunk 在向量索引中的节点 ID
+    index_node_hash: str  # hash of sign_content
+    hit_count: int
+    enabled: bool
+    status: str  # 'completed'
+    created_at: int  # timestamp
+    updated_at: int  # timestamp
+    completed_at: int  # timestamp
+    created_by: str  # user id
+    child_chunks: list
+    error: Any | None
+    stopped_at: int | None  # timestamp
+    disabled_at: int | None  # timestamp

davidkhala/ai/agent/dify/{ops/console → console}/__init__.py

@@ -1,9 +1,15 @@
+from enum import Enum
+
 from davidkhala.utils.http_request import Request
 
 
 class API(Request):
+
     def __init__(self, base_url='http://localhost'):
+        """
+        :param base_url: "{protocol}://{host}". For Dify cloud, it is 'https://cloud.dify.ai'
+        """
         super().__init__()
         self.base_url = f"{base_url}/console/api"
-        self.__enter__()
+        self.open()
 

davidkhala/ai/agent/dify/console/knowledge/dataset.py

@@ -0,0 +1,61 @@
+from time import sleep
+
+from davidkhala.ai.agent.dify.console.session import ConsoleDerived
+from davidkhala.ai.agent.dify.const import IndexingStatus
+from davidkhala.ai.agent.dify.interface import IndexingError
+
+
+class Operation(ConsoleDerived):
+    def website_sync(self, dataset: str, document: str, *, wait_until=True):
+        """
+        cannot be used towards a pipeline dataset. Otherwise, you will see error "no website import info found"
+        """
+        doc_url = f"{self.base_url}/datasets/{dataset}/documents/{document}"
+
+        r = self.request(f"{doc_url}/website-sync", "GET")
+        assert r == {"result": "success"}
+        if wait_until:
+            return self.wait_until(dataset, document)
+        return None
+
+    def retry(self, dataset: str, *documents: str, wait_until=True):
+        """
+        It cannot trigger rerun on success documents
+        """
+        url = f"{self.base_url}/datasets/{dataset}/retry"
+        self.request(url, "POST", json={
+            'document_ids': documents,
+        })
+        # response status code will be 204
+        if wait_until:
+            return [self.wait_until(dataset, document) for document in documents]
+        return None
+
+    def rerun(self, dataset: str, *documents: str):
+        for document in documents:
+            try:
+                self.website_sync(dataset, document)
+                assert False, "expect IndexingError"
+            except IndexingError:
+                pass
+        return self.retry(dataset, *documents)
+
+    def wait_until(self, dataset: str, document: str, *,
+                   expect_status=None,
+                   from_status=None,
+                   interval=1
+                   ):
+        if not expect_status:
+            expect_status = [IndexingStatus.FAILED, IndexingStatus.COMPLETED]
+        url = f"{self.base_url}/datasets/{dataset}/documents/{document}/indexing-status"
+        if from_status is None:
+            from_status = [IndexingStatus.WAITING, IndexingStatus.PARSING]
+        r = self.request(url, "GET")
+        status = r['indexing_status']
+        assert status in from_status, f"current status: {status}, expect: {from_status}"
+        while status not in expect_status:
+            sleep(interval)
+            r = self.request(url, "GET")
+            status = r['indexing_status']
+        if status == IndexingStatus.FAILED: raise IndexingError(r['error'])
+        return r

davidkhala/ai/agent/dify/console/knowledge/pipeline.py

@@ -0,0 +1,127 @@
+from typing import Any
+
+from davidkhala.utils.http_request.stream import Request as StreamRequest, as_sse
+from pydantic import BaseModel, Field
+
+from davidkhala.ai.agent.dify.console.session import ConsoleDerived
+from davidkhala.ai.agent.dify.model import User
+from davidkhala.ai.agent.dify.model.knowledge import Dataset, Document
+from davidkhala.ai.agent.dify.model.workflow import NodeProtocol, Graph
+
+
+class RAGPipelineVariable(BaseModel):
+    label: str
+    variable: str
+    type: str
+    belong_to_node_id: str
+    max_length: int | None = None
+    required: bool = False
+    unit: str | None = None
+    default_value: Any | None = None
+    options: list[Any] = Field(default_factory=list)
+    placeholder: str | None = None
+    tooltips: str | None = None
+    allowed_file_types: str | None = None
+    allow_file_extension: str | None = None
+    allow_file_upload_methods: str | None = None
+
+
+class PipelineModel(BaseModel):
+    id: str
+    graph: Graph
+    features: dict[str, Any] = Field(default_factory=dict)
+    hash: str
+    version: str
+    marked_name: str = ""
+    marked_comment: str = ""
+    created_by: User
+    created_at: int
+    updated_by: User | None = None
+    updated_at: int
+    tool_published: bool = False
+    environment_variables: list[dict[str, Any]]
+    conversation_variables: list[dict[str, Any]]
+    rag_pipeline_variables: list[RAGPipelineVariable]
+
+
+class DatasetResult(Dataset):
+    chunk_structure: str
+
+
+class RunResult(BaseModel):
+    batch: str
+    dataset: DatasetResult
+    documents: list[Document]
+
+
+class Pipeline(ConsoleDerived):
+
+    def async_run(self, pipeline: str, node: NodeProtocol, inputs: dict, datasource_info_list: list[dict]) -> RunResult:
+        """Ingest new document"""
+        url = f"{self.base_url}/rag/pipelines/{pipeline}/workflows/published/run"
+        r = self.request(url, "POST", json={
+            'inputs': inputs,
+            'start_node_id': node.id,
+            'is_preview': False,
+            'response_mode': "blocking",
+            "datasource_info_list": datasource_info_list,
+            'datasource_type': node.datasource_type
+        })
+        return RunResult.model_validate(r)
+
+    def get(self, pipeline: str):
+        url = f"{self.base_url}/rag/pipelines/{pipeline}/workflows/publish"
+        r = self.request(url, "GET")
+        return PipelineModel.model_validate(r)
+
+
+class Datasource(ConsoleDerived):
+    class FirecrawlOutput(BaseModel):
+        source_url: str
+        description: str
+        title: str
+        credential_id: str
+        content: str
+
+    def run_firecrawl(self, pipeline: str, node: NodeProtocol,
+                      *,
+                      inputs: dict,
+                      credential_id: str
+                      ):
+
+        url = f"{self.base_url}/rag/pipelines/{pipeline}/workflows/published/datasource/nodes/{node.id}/run"
+
+        stream_request = StreamRequest(self)
+        response = stream_request.request(url, 'POST', json={
+            'inputs': inputs,
+            'datasource_type': node.datasource_type,
+            'credential_id': credential_id,
+            "response_mode": "streaming"
+        })
+
+        for data in as_sse(response):
+            event = data['event']
+            if event == 'datasource_completed':
+                return data['data']
+            else:
+                assert event == 'datasource_processing'
+                print(data)
+        return None
+
+    def upload(self):
+        "http://localhost/console/api/files/upload?source=datasets"
+        # TODO
+        "form data"
+        {
+            "file": "body"
+        }
+        r = {
+            "id": "3898db5b-eb72-4f11-b507-628ad5d28887",
+            "name": "Professional Diploma Meister Power Electrical Engineering - Technological and Higher Education Institute of Hong Kong.html",
+            "size": 254362,
+            "extension": "html",
+            "mime_type": "text\/html",
+            "created_by": "dbd0b38b-5ef1-4123-8c3f-0c82eb1feacd",
+            "created_at": 1764943811,
+            "source_url": "\/files\/3898db5b-eb72-4f11-b507-628ad5d28887\/file-preview?timestamp=1764943811&nonce=43b0ff5a13372415be79de4cc7ef398c&sign=7OJ2wiVYc4tygl7yvM1sPn7s0WXDlhHxgX76bsGTD94%3D"
+        }

davidkhala/ai/agent/dify/{ops/console → console}/plugin.py

@@ -1,15 +1,12 @@
 from time import sleep
 
-from davidkhala.ai.agent.dify.ops.console import API
-from davidkhala.ai.agent.dify.ops.console.session import ConsoleUser
+from davidkhala.ai.agent.dify.console.session import ConsoleUser, ConsoleDerived
 
 
-class ConsolePlugin(API):
+class ConsolePlugin(ConsoleDerived):
     def __init__(self, context: ConsoleUser):
-        super().__init__()
+        super().__init__(context)
         self.base_url = f"{context.base_url}/workspaces/current/plugin"
-        self.session.cookies = context.session.cookies
-        self.options = context.options
 
     def upgrade(self, *plugin_names: str) -> list[dict]:
         versions = self.latest_version(*plugin_names)
@@ -40,7 +37,7 @@ class ConsolePlugin(API):
         return _
 
     def get(self, *plugin_names: str) -> list[dict]:
-        "inspect installed plugins"
+        """inspect installed plugins"""
         url = f"{self.base_url}/list/installations/ids"
         r = self.request(url, method="POST", json={
             'plugin_ids': plugin_names,
@@ -66,3 +63,20 @@ class ConsolePlugin(API):
         for name in plugin_names:
             r = self.get(name)
             self.uninstall(r[0]['id'])
+
+
+class ConsoleTool(ConsoleDerived):
+    def __init__(self, context: ConsoleUser):
+        super().__init__(context)
+        self.base_url = f"{context.base_url}/workspaces/current/tool-provider"
+
+    def credential_id_by(self, name, owner: str, plugin: str):
+        for c in self.credentials_of(owner, plugin):
+            if c['name'] == name:
+                return c['id']
+        return None
+
+    def credentials_of(self, owner: str, plugin: str):
+        url = f"{self.base_url}/builtin/{owner}/{plugin}/{plugin}/credential/info"
+        r = self.request(url, method="GET")
+        return r['credentials']

davidkhala/ai/agent/dify/console/session.py

@@ -0,0 +1,50 @@
+from base64 import b64encode
+
+from requests.cookies import RequestsCookieJar
+
+from davidkhala.ai.agent.dify.console import API
+
+
+class ConsoleUser(API):
+    def login(self, email, password,
+              *,
+              remember_me=True,
+              language="en-US"
+              ) -> RequestsCookieJar:
+        url = f"{self.base_url}/login"
+
+        r = self.request(url, "POST", json={
+            'email': email,
+            'password': b64encode(password.encode()).decode(),  # use base64 from dify 1.11
+            'remember_me': remember_me,
+            'language': language,
+        })
+        assert r == {"result": "success"}
+        self.options['headers']['x-csrf-token'] = self.session.cookies.get("csrf_token")
+        return self.session.cookies
+
+    def set_tokens(self, *, csrf, access):
+        """workaround for federated login"""
+        self.session.cookies.set(name="__Host-csrf_token", value=csrf)
+        self.session.cookies.set(name="__Host-access_token", value=access)
+
+
+        self.options['headers']['x-csrf-token'] = csrf
+
+    @property
+    def me(self) -> dict:
+        url = f"{self.base_url}/account/profile"
+        return self.request(url, "GET")
+
+    @property
+    def workspace(self) -> dict:
+        url = f"{self.base_url}/features"
+        return self.request(url, "GET")
+
+
+class ConsoleDerived(API):
+    def __init__(self, context: ConsoleUser):
+        super().__init__()
+        self.base_url = context.base_url
+        self.session.cookies = context.session.cookies
+        self.options = context.options