chore: import upstream snapshot with attribution

sdk/python/ragflow_sdk/modules/__init__.py

@@ -0,0 +1,15 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#

sdk/python/ragflow_sdk/modules/agent.py

@@ -0,0 +1,69 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from .base import Base
+from .session import Session
+
+
+class Agent(Base):
+    def __init__(self, rag, res_dict):
+        self.id = None
+        self.avatar = None
+        self.canvas_type = None
+        self.description = None
+        self.dsl = None
+        super().__init__(rag, res_dict)
+
+    class Dsl(Base):
+        def __init__(self, rag, res_dict):
+            self.answer = []
+            self.components = {"begin": {"downstream": ["Answer:China"], "obj": {"component_name": "Begin", "params": {}}, "upstream": []}}
+            self.graph = {
+                "edges": [],
+                "nodes": [{"data": {"label": "Begin", "name": "begin"}, "id": "begin", "position": {"x": 50, "y": 200}, "sourcePosition": "left", "targetPosition": "right", "type": "beginNode"}],
+            }
+            self.history = []
+            self.messages = []
+            self.path = []
+            self.reference = []
+            super().__init__(rag, res_dict)
+
+    def create_session(self, **kwargs) -> Session:
+        res = self.post(f"/agents/{self.id}/sessions", json=kwargs)
+        res = res.json()
+        if res.get("code") == 0:
+            return Session(self.rag, res.get("data"))
+        raise Exception(res.get("message"))
+
+    def list_sessions(self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str = None) -> list[Session]:
+        res = self.get(f"/agents/{self.id}/sessions", {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id})
+        res = res.json()
+        if res.get("code") == 0:
+            result_list = []
+            for data in res.get("data"):
+                temp_agent = Session(self.rag, data)
+                result_list.append(temp_agent)
+            return result_list
+        raise Exception(res.get("message"))
+
+    def delete_sessions(self, ids: list[str] | None = None, delete_all: bool = False):
+        payload = {"ids": ids}
+        if delete_all:
+            payload["delete_all"] = True
+        res = self.rm(f"/agents/{self.id}/sessions", payload)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))

sdk/python/ragflow_sdk/modules/base.py

@@ -0,0 +1,62 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+
+class Base:
+    def __init__(self, rag, res_dict):
+        self.rag = rag
+        self._update_from_dict(rag, res_dict)
+
+    def _update_from_dict(self, rag, res_dict):
+        for k, v in res_dict.items():
+            if isinstance(v, dict):
+                self.__dict__[k] = Base(rag, v)
+            else:
+                self.__dict__[k] = v
+
+    def to_json(self):
+        pr = {}
+        for name in dir(self):
+            value = getattr(self, name)
+            if not name.startswith("__") and not callable(value) and name != "rag":
+                if isinstance(value, Base):
+                    pr[name] = value.to_json()
+                else:
+                    pr[name] = value
+        return pr
+
+    def post(self, path, json=None, stream=False, files=None):
+        res = self.rag.post(path, json, stream=stream, files=files)
+        return res
+
+    def get(self, path, params=None):
+        res = self.rag.get(path, params)
+        return res
+
+    def rm(self, path, json):
+        res = self.rag.delete(path, json)
+        return res
+
+    def put(self, path, json):
+        res = self.rag.put(path, json)
+        return res
+
+    def patch(self, path, json):
+        res = self.rag.patch(path, json)
+        return res
+
+    def __str__(self):
+        return str(self.to_json())

sdk/python/ragflow_sdk/modules/chat.py

@@ -0,0 +1,67 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+
+from .base import Base
+from .session import Session
+
+
+class Chat(Base):
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.name = "assistant"
+        self.icon = ""
+        self.dataset_ids = []
+        self.llm_id = None
+        self.llm_setting = {}
+        self.prompt_config = {}
+        self.similarity_threshold = 0.2
+        self.vector_similarity_weight = 0.3
+        self.top_n = 6
+        self.top_k = 1024
+        self.rerank_id = ""
+        super().__init__(rag, res_dict)
+
+    def update(self, update_message: dict):
+        if not isinstance(update_message, dict):
+            raise Exception("ValueError('`update_message` must be a dict')")
+        res = self.patch(f"/chats/{self.id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+
+    def create_session(self, name: str = "New session") -> Session:
+        res = self.post(f"/chats/{self.id}/sessions", {"name": name})
+        res = res.json()
+        if res.get("code") == 0:
+            return Session(self.rag, res["data"])
+        raise Exception(res.get("message"))
+
+    def list_sessions(self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str = None, name: str = None, user_id: str = None) -> list[Session]:
+        res = self.get(f"/chats/{self.id}/sessions", {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id, "name": name, "user_id": user_id})
+        res = res.json()
+        if res.get("code") == 0:
+            result_list = []
+            for data in res["data"]:
+                result_list.append(Session(self.rag, data))
+            return result_list
+        raise Exception(res["message"])
+
+    def delete_sessions(self, ids: list[str] | None = None, delete_all: bool = False):
+        res = self.rm(f"/chats/{self.id}/sessions", {"ids": ids, "delete_all": delete_all})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))

sdk/python/ragflow_sdk/modules/chunk.py

@@ -0,0 +1,61 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from .base import Base
+
+
+class ChunkUpdateError(Exception):
+    def __init__(self, code=None, message=None, details=None):
+        self.code = code
+        self.message = message
+        self.details = details
+        super().__init__(message)
+
+
+class Chunk(Base):
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.content = ""
+        self.important_keywords = []
+        self.tag_kwd = []
+        self.questions = []
+        self.create_time = ""
+        self.create_timestamp = 0.0
+        self.dataset_id = None
+        self.document_name = ""
+        self.document_keyword = ""
+        self.document_id = ""
+        self.available = True
+        # Additional fields for retrieval results
+        self.similarity = 0.0
+        self.vector_similarity = 0.0
+        self.term_similarity = 0.0
+        self.positions = []
+        self.doc_type = ""
+        for k in list(res_dict.keys()):
+            if k not in self.__dict__:
+                res_dict.pop(k)
+        super().__init__(rag, res_dict)
+
+        # for backward compatibility
+        if not self.document_name:
+            self.document_name = self.document_keyword
+
+    def update(self, update_message: dict):
+        res = self.patch(f"/datasets/{self.dataset_id}/documents/{self.document_id}/chunks/{self.id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise ChunkUpdateError(code=res.get("code"), message=res.get("message"), details=res.get("details"))

sdk/python/ragflow_sdk/modules/dataset.py

@@ -0,0 +1,183 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+from typing import Any
+
+from .base import Base
+from .document import Document
+
+
+class DataSet(Base):
+    class ParserConfig(Base):
+        def __init__(self, rag, res_dict):
+            super().__init__(rag, res_dict)
+
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.name = ""
+        self.avatar = ""
+        self.tenant_id = None
+        self.description = ""
+        self.embedding_model = ""
+        self.permission = "me"
+        self.document_count = 0
+        self.chunk_count = 0
+        self.chunk_method = "naive"
+        self.parser_config = None
+        self.pagerank = 0
+        for k in list(res_dict.keys()):
+            if k not in self.__dict__:
+                res_dict.pop(k)
+        super().__init__(rag, res_dict)
+
+    def update(self, update_message: dict):
+        res = self.put(f"/datasets/{self.id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+
+        self._update_from_dict(self.rag, res.get("data", {}))
+        return self
+
+    def upload_documents(self, document_list: list[dict]):
+        url = f"/datasets/{self.id}/documents"
+        files = [("file", (ele["display_name"], ele["blob"])) for ele in document_list]
+        res = self.post(path=url, json=None, files=files)
+        res = res.json()
+        if res.get("code") == 0:
+            doc_list = []
+            for doc in res["data"]:
+                document = Document(self.rag, doc)
+                doc_list.append(document)
+            return doc_list
+        raise Exception(res.get("message"))
+
+    def list_documents(
+        self,
+        id: str | None = None,
+        ids: list[str] | None = None,
+        name: str | None = None,
+        keywords: str | None = None,
+        page: int = 1,
+        page_size: int = 30,
+        orderby: str = "create_time",
+        desc: bool = True,
+        create_time_from: int = 0,
+        create_time_to: int = 0,
+    ):
+        # Validate that id and ids are not used together
+        if id and ids:
+            raise ValueError("Cannot use both 'id' and 'ids' parameters at the same time.")
+
+        params = {
+            "id": id,
+            "name": name,
+            "keywords": keywords,
+            "page": page,
+            "page_size": page_size,
+            "orderby": orderby,
+            "desc": desc,
+            "create_time_from": create_time_from,
+            "create_time_to": create_time_to,
+        }
+        # Handle ids parameter - convert to multiple query params
+        if ids:
+            for doc_id in ids:
+                params.append(("ids", doc_id))
+        res = self.get(f"/datasets/{self.id}/documents", params=params)
+        res = res.json()
+        documents = []
+        if res.get("code") == 0:
+            for document in res["data"].get("docs"):
+                documents.append(Document(self.rag, document))
+            return documents
+        raise Exception(res["message"])
+
+    def delete_documents(self, ids: list[str] | None = None, delete_all: bool = False):
+        res = self.rm(f"/datasets/{self.id}/documents", {"ids": ids, "delete_all": delete_all})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+
+    def _get_documents_status(self, document_ids):
+        import time
+
+        terminal_states = {"DONE", "FAIL", "CANCEL"}
+        interval_sec = 1
+        pending = set(document_ids)
+        finished = []
+        while pending:
+            for doc_id in list(pending):
+
+                def fetch_doc(doc_id: str) -> Document | None:
+                    try:
+                        docs = self.list_documents(id=doc_id)
+                        return docs[0] if docs else None
+                    except Exception:
+                        return None
+
+                doc = fetch_doc(doc_id)
+                if doc is None:
+                    continue
+                if isinstance(doc.run, str) and doc.run.upper() in terminal_states:
+                    finished.append((doc_id, doc.run, doc.chunk_count, doc.token_count))
+                    pending.discard(doc_id)
+                elif float(doc.progress or 0.0) >= 1.0:
+                    finished.append((doc_id, "DONE", doc.chunk_count, doc.token_count))
+                    pending.discard(doc_id)
+            if pending:
+                time.sleep(interval_sec)
+        return finished
+
+    def async_parse_documents(self, document_ids):
+        res = self.post(f"/datasets/{self.id}/chunks", {"document_ids": document_ids})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))
+
+    def parse_documents(self, document_ids):
+        try:
+            self.async_parse_documents(document_ids)
+            self._get_documents_status(document_ids)
+        except KeyboardInterrupt:
+            self.async_cancel_parse_documents(document_ids)
+
+        return self._get_documents_status(document_ids)
+
+    def async_cancel_parse_documents(self, document_ids):
+        res = self.rm(f"/datasets/{self.id}/chunks", {"document_ids": document_ids})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))
+
+    def get_auto_metadata(self) -> dict[str, Any]:
+        """
+        Retrieve auto-metadata configuration for a dataset via SDK.
+        """
+        res = self.get(f"/datasets/{self.id}/metadata/config")
+        res = res.json()
+        if res.get("code") == 0:
+            return res["data"]
+        raise Exception(res["message"])
+
+    def update_auto_metadata(self, **config: Any) -> dict[str, Any]:
+        """
+        Update auto-metadata configuration for a dataset via SDK.
+        """
+        res = self.put(f"/datasets/{self.id}/metadata/config", config)
+        res = res.json()
+        if res.get("code") == 0:
+            return res["data"]
+        raise Exception(res["message"])

sdk/python/ragflow_sdk/modules/document.py

@@ -0,0 +1,104 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import json
+
+from .base import Base
+from .chunk import Chunk
+
+
+class Document(Base):
+    class ParserConfig(Base):
+        def __init__(self, rag, res_dict):
+            super().__init__(rag, res_dict)
+
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.name = ""
+        self.thumbnail = None
+        self.dataset_id = None
+        self.chunk_method = "naive"
+        self.parser_config = {"pages": [[1, 1000000]]}
+        self.source_type = "local"
+        self.type = ""
+        self.created_by = ""
+        self.size = 0
+        self.token_count = 0
+        self.chunk_count = 0
+        self.progress = 0.0
+        self.progress_msg = ""
+        self.process_begin_at = None
+        self.process_duration = 0.0
+        self.run = "0"
+        self.status = "1"
+        self.meta_fields = {}
+        for k in list(res_dict.keys()):
+            if k not in self.__dict__:
+                res_dict.pop(k)
+        super().__init__(rag, res_dict)
+
+    def update(self, update_message: dict):
+        if "meta_fields" in update_message:
+            if not isinstance(update_message["meta_fields"], dict):
+                raise Exception("meta_fields must be a dictionary")
+        res = self.patch(f"/datasets/{self.dataset_id}/documents/{self.id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+
+        self._update_from_dict(self.rag, res.get("data", {}))
+        return self
+
+    def download(self):
+        res = self.get(f"/datasets/{self.dataset_id}/documents/{self.id}")
+        error_keys = set(["code", "message"])
+        try:
+            response = res.json()
+            actual_keys = set(response.keys())
+            if actual_keys == error_keys:
+                raise Exception(response.get("message"))
+            else:
+                return res.content
+        except json.JSONDecodeError:
+            return res.content
+
+    def list_chunks(self, page=1, page_size=30, keywords="", id=""):
+        data = {"keywords": keywords, "page": page, "page_size": page_size, "id": id}
+        res = self.get(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", data)
+        res = res.json()
+        if res.get("code") == 0:
+            chunks = []
+            for data in res["data"].get("chunks"):
+                chunk = Chunk(self.rag, data)
+                chunks.append(chunk)
+            return chunks
+        raise Exception(res.get("message"))
+
+    def add_chunk(self, content: str, important_keywords: list[str] = [], questions: list[str] = [], image_base64: str | None = None, *, tag_kwd: list[str] = []):
+        body = {"content": content, "important_keywords": important_keywords, "tag_kwd": tag_kwd, "questions": questions}
+        if image_base64 is not None:
+            body["image_base64"] = image_base64
+        res = self.post(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", body)
+        res = res.json()
+        if res.get("code") == 0:
+            return Chunk(self.rag, res["data"].get("chunk"))
+        raise Exception(res.get("message"))
+
+    def delete_chunks(self, ids: list[str] | None = None, delete_all: bool = False):
+        res = self.rm(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", {"chunk_ids": ids, "delete_all": delete_all})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))

sdk/python/ragflow_sdk/modules/memory.py

@@ -0,0 +1,87 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from .base import Base
+
+
+class Memory(Base):
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.name = ""
+        self.avatar = None
+        self.tenant_id = None
+        self.owner_name = ""
+        self.memory_type = ["raw"]
+        self.storage_type = "table"
+        self.embd_id = ""
+        self.llm_id = ""
+        self.permissions = "me"
+        self.description = ""
+        self.memory_size = 5 * 1024 * 1024
+        self.forgetting_policy = "FIFO"
+        self.temperature = (0.5,)
+        self.system_prompt = ""
+        self.user_prompt = ""
+        for k in list(res_dict.keys()):
+            if k not in self.__dict__:
+                res_dict.pop(k)
+        super().__init__(rag, res_dict)
+
+    def update(self, update_dict: dict):
+        res = self.put(f"/memories/{self.id}", update_dict)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        self._update_from_dict(self.rag, res.get("data", {}))
+        return self
+
+    def get_config(self):
+        res = self.get(f"/memories/{self.id}/config")
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        self._update_from_dict(self.rag, res.get("data", {}))
+        return self
+
+    def list_memory_messages(self, agent_id: str | list[str] = None, keywords: str = None, page: int = 1, page_size: int = 50):
+        params = {"agent_id": agent_id, "keywords": keywords, "page": page, "page_size": page_size}
+        res = self.get(f"/memories/{self.id}", params)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        return res["data"]
+
+    def forget_message(self, message_id: int):
+        res = self.rm(f"/messages/{self.id}:{message_id}", {})
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        return True
+
+    def update_message_status(self, message_id: int, status: bool):
+        update_message = {"status": status}
+        res = self.put(f"/messages/{self.id}:{message_id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        return True
+
+    def get_message_content(self, message_id: int) -> dict:
+        res = self.get(f"/messages/{self.id}:{message_id}/content")
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res["message"])
+        return res["data"]

sdk/python/ragflow_sdk/modules/session.py

@@ -0,0 +1,186 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import json
+import logging
+
+from .base import Base
+
+logger = logging.getLogger(__name__)
+
+
+class Session(Base):
+    def __init__(self, rag, res_dict):
+        self.id = None
+        self.name = "New session"
+        self.messages = [{"role": "assistant", "content": "Hi! I am your assistant, can I help you?"}]
+        for key, value in res_dict.items():
+            if key == "chat_id" and value is not None:
+                self.chat_id = None
+                self.__session_type = "chat"
+            if key == "agent_id" and value is not None:
+                self.agent_id = None
+                self.__session_type = "agent"
+        super().__init__(rag, res_dict)
+
+    def ask(
+        self,
+        question="",
+        stream=False,
+        inputs=None,
+        release=None,
+        return_trace=None,
+        **kwargs,
+    ):
+        """
+        Ask a question to the session.
+
+        Parameters
+        ----------
+        question : str
+            The user's question. May be empty when the agent is driven solely by
+            Begin component inputs.
+        stream : bool
+            If ``True``, yields ``Message`` objects as they arrive (SSE streaming).
+            If ``False``, yields a single ``Message`` with the final answer.
+        inputs : dict, optional
+            Values for variables declared on the agent's **Begin** component. Each
+            value must be a dict containing at least a ``"value"`` key, and may
+            include ``"type"``. Example::
+
+                session.ask(
+                    "",
+                    stream=False,
+                    inputs={"key1": {"type": "line", "value": "hello"}},
+                )
+
+            Only meaningful for agent sessions; ignored for chat sessions.
+        release : bool, optional
+            If ``True``, run against the latest published agent version instead of
+            the editable draft. Only meaningful for agent sessions.
+        return_trace : bool, optional
+            If ``True``, include execution trace information in the response.
+            Only meaningful for agent sessions.
+        **kwargs
+            Additional fields forwarded verbatim to the completion endpoint
+            (e.g. ``session_id``, ``files``, ``user_id``, ``custom_header``).
+            See the HTTP API reference for the full list.
+        """
+        if inputs is not None:
+            kwargs["inputs"] = inputs
+        if release is not None:
+            kwargs["release"] = release
+        if return_trace is not None:
+            kwargs["return_trace"] = return_trace
+
+        if inputs is not None or release is not None or return_trace is not None:
+            logger.debug(
+                "Session.ask explicit-params session_type=%s session_id=%s input_keys=%s release=%s return_trace=%s",
+                self.__session_type,
+                getattr(self, "id", None),
+                list(inputs.keys()) if isinstance(inputs, dict) else None,
+                release,
+                return_trace,
+            )
+
+        if self.__session_type == "agent":
+            res = self._ask_agent(question, stream, **kwargs)
+        elif self.__session_type == "chat":
+            res = self._ask_chat(question, stream, **kwargs)
+        else:
+            raise Exception(f"Unknown session type: {self.__session_type}")
+
+        if stream:
+            for line in res.iter_lines(decode_unicode=True):
+                if not line:
+                    continue  # Skip empty lines
+                line = line.strip()
+                if line.startswith("data:"):
+                    content = line[len("data:") :].strip()
+                    if content == "[DONE]":
+                        break  # End of stream
+                else:
+                    content = line
+
+                try:
+                    json_data = json.loads(content)
+                except json.JSONDecodeError:
+                    continue  # Skip lines that are not valid JSON
+
+                event = json_data.get("event", None)
+                if event and event != "message":
+                    continue
+
+                if (self.__session_type == "agent" and event == "message_end") or (self.__session_type == "chat" and json_data.get("data") is True):
+                    return
+                if self.__session_type == "agent":
+                    yield self._structure_answer(json_data)
+                else:
+                    yield self._structure_answer(json_data["data"])
+        else:
+            try:
+                json_data = res.json()
+            except ValueError:
+                raise Exception(f"Invalid response {res}")
+            yield self._structure_answer(json_data["data"])
+
+    def _structure_answer(self, json_data):
+        answer = ""
+        if self.__session_type == "agent":
+            answer = json_data["data"]["content"]
+        elif self.__session_type == "chat":
+            answer = json_data["answer"]
+        reference = json_data.get("reference", {})
+        temp_dict = {"content": answer, "role": "assistant"}
+        if reference and "chunks" in reference:
+            chunks = reference["chunks"]
+            temp_dict["reference"] = chunks
+        message = Message(self.rag, temp_dict)
+        return message
+
+    def _ask_chat(self, question: str, stream: bool, **kwargs):
+        json_data = {"question": question, "stream": stream, "session_id": self.id}
+        json_data.update(kwargs)
+        res = self.post(f"/chats/{self.chat_id}/completions", json_data, stream=stream)
+        return res
+
+    def _ask_agent(self, question: str, stream: bool, **kwargs):
+        json_data = {
+            "agent_id": self.agent_id,
+            "query": question,
+            "stream": stream,
+            "session_id": self.id,
+            "openai-compatible": False,
+        }
+        json_data.update(kwargs)
+        res = self.post("/agents/chat/completions", json_data, stream=stream)
+        return res
+
+    def update(self, update_message):
+        res = self.patch(f"/chats/{self.chat_id}/sessions/{self.id}", update_message)
+        res = res.json()
+        if res.get("code") != 0:
+            raise Exception(res.get("message"))
+
+
+class Message(Base):
+    def __init__(self, rag, res_dict):
+        self.content = "Hi! I am your assistant, can I help you?"
+        self.reference = None
+        self.role = "assistant"
+        self.prompt = None
+        self.id = None
+        super().__init__(rag, res_dict)