feat: add PageIndex SDK with local/cloud dual-mode support (#207)

pageindex/collection.py

@@ -0,0 +1,69 @@
+# pageindex/collection.py
+from __future__ import annotations
+from typing import AsyncIterator
+from .events import QueryEvent
+from .backend.protocol import Backend
+
+
+class QueryStream:
+    """Wraps backend.query_stream() as an async iterable object."""
+
+    def __init__(self, backend: Backend, collection: str, question: str,
+                 doc_ids: list[str] | None = None):
+        self._backend = backend
+        self._collection = collection
+        self._question = question
+        self._doc_ids = doc_ids
+
+    async def stream_events(self) -> AsyncIterator[QueryEvent]:
+        async for event in self._backend.query_stream(
+            self._collection, self._question, self._doc_ids
+        ):
+            yield event
+
+    def __aiter__(self):
+        return self.stream_events()
+
+
+class Collection:
+    def __init__(self, name: str, backend: Backend):
+        self._name = name
+        self._backend = backend
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def add(self, file_path: str) -> str:
+        return self._backend.add_document(self._name, file_path)
+
+    def list_documents(self) -> list[dict]:
+        return self._backend.list_documents(self._name)
+
+    def get_document(self, doc_id: str, include_text: bool = False) -> dict:
+        return self._backend.get_document(self._name, doc_id, include_text=include_text)
+
+    def get_document_structure(self, doc_id: str) -> list:
+        return self._backend.get_document_structure(self._name, doc_id)
+
+    def get_page_content(self, doc_id: str, pages: str) -> list:
+        return self._backend.get_page_content(self._name, doc_id, pages)
+
+    def delete_document(self, doc_id: str) -> None:
+        self._backend.delete_document(self._name, doc_id)
+
+    def query(self, question: str, doc_ids: list[str] | None = None,
+              stream: bool = False) -> str | QueryStream:
+        """Query documents in this collection.
+
+        - stream=False: returns answer string (sync)
+        - stream=True: returns async iterable of QueryEvent
+
+        Usage:
+            answer = col.query("question")
+            async for event in col.query("question", stream=True):
+                ...
+        """
+        if stream:
+            return QueryStream(self._backend, self._name, question, doc_ids)
+        return self._backend.query(self._name, question, doc_ids)