Merge pull request #15 from Querent-ai/collector_interfaces

add some collector interfaces

Author: saraswatpuneet

querent/connectors/__init__.py renamed to querent/collectors/__init__.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator
+
+from querent.collectors.collector_result import CollectorResult
+
+class Collector(ABC):
+    @abstractmethod
+    async def connect(self):
+        pass
+
+    @abstractmethod
+    async def poll(self) -> AsyncGenerator[CollectorResult, None]:
+        pass
+
+    @abstractmethod
+    async def disconnect(self):
+        pass

querent/collectors/collector_base.py

@@ -0,0 +1,34 @@
+from typing import List, Dict
+from pathlib import Path
+from enum import Enum
+
+class CollectorErrorKind(Enum):
+    NotFound = "not_found"
+    Unauthorized = "unauthorized"
+    incompatible = "incompatible"
+    NotSupported = "not_supported"
+
+class CollectorResolverError(Exception):
+    def __init__(self, kind: CollectorErrorKind, message: str):
+        super().__init__(message)
+        self.kind = kind
+
+
+class CollectorError(Exception):
+    def __init__(self, kind: CollectorErrorKind, source: any):
+        self.kind = kind
+        self.source = source
+
+
+class NotFoundError(CollectorError):
+    def __init__(self, message: str):
+        super().__init__(CollectorErrorKind.NotFound, message)
+
+class UnauthorizedError(CollectorError):
+    def __init__(self, message: str):
+        super().__init__(CollectorErrorKind.unauthorized, message)
+
+
+class IncompatibleError(CollectorError):
+    def __init__(self, message: str):
+        super().__init__(CollectorErrorKind.incompatible, message)

querent/collectors/collector_errors.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+from enum import Enum
+from querent.collectors.collector_base import Collector
+from querent.collectors.collector_errors import CollectorResolverError, CollectorErrorKind
+from querent.config.collector_config import CollectorBackend
+
+
+class CollectorFactory(ABC):
+    @abstractmethod
+    def backend(self) -> CollectorBackend:
+        pass
+
+    @abstractmethod
+    async def resolve(self, uri: str) -> Optional[CollectorBackend]:
+        pass
+
+
+class UnsupportedCollector(CollectorFactory):
+    def __init__(self, backend: CollectorBackend, message: str):
+        self.backend = backend
+        self.message = message
+
+    async def resolve(self, uri: str) -> Optional[Collector]:
+        raise CollectorResolverError(
+            CollectorErrorKind.NotSupported, self.backend, self.message
+        )

querent/collectors/collector_factory.py

@@ -0,0 +1,32 @@
+from typing import Optional
+from querent.collectors.fs.fs_collector import FSCollectorFactory
+from querent.config.collector_config import CollectorBackend
+from querent.collectors.collector_base import Collector
+from querent.collectors.collector_errors import CollectorResolverError, CollectorErrorKind
+from querent.common.uri import Protocol, Uri
+
+class CollectorResolver:
+    def __init__(self):
+        self.collector_factories = {
+            CollectorBackend.LocalFile: FSCollectorFactory(),
+            # Add other collector factories as needed
+        }
+
+    def resolve(self, uri: Uri) -> Optional[Collector]:
+        backend = self._determine_backend(uri.protocol)
+        
+        if backend in self.collector_factories:
+            factory = self.collector_factories[backend]
+            return factory.resolve(uri)
+        else:
+            raise CollectorResolverError(
+                CollectorErrorKind.NotSupported, backend, "Unsupported backend"
+            )
+
+    def _determine_backend(self, protocol: Protocol) -> CollectorBackend:
+        if protocol.is_file_storage():
+            return CollectorBackend.LocalFile
+        else:
+            raise CollectorResolverError(
+                CollectorErrorKind.NotSupported, "Unknown backend", "Unknown backend"
+            )

querent/collectors/collector_resolver.py

@@ -0,0 +1,56 @@
+from typing import Any, Union
+
+class CollectorResult:
+    def __init__(self, value: Any, error: str = None) -> None:
+        self.value = value
+        self.error = error
+
+    def __str__(self):
+        if self.error:
+            return f"Error: {self.error}"
+        return f"Value: {self.value}"
+
+    def is_error(self) -> bool:
+        return self.error is not None
+
+    @classmethod
+    def success(cls, value: Any) -> "CollectorResult":
+        return cls(value)
+
+    @classmethod
+    def error(cls, error: str) -> "CollectorResult":
+        return cls(None, error)
+
+    def unwrap(self) -> Any:
+        if self.error:
+            raise ValueError(self.error)
+        return self.value
+
+    def unwrap_or(self, default: Any) -> Any:
+        return self.value if not self.error else default
+
+    def __eq__(self, other: Union[Any, "CollectorResult"]) -> bool:
+        if isinstance(other, CollectorResult):
+            return self.value == other.value and self.error == other.error
+        return self.value == other
+
+    def __hash__(self) -> int:
+        return hash((self.value, self.error))
+
+
+# Usage example
+if __name__ == "__main__":
+    success_result = CollectorResult.success("Success value")
+    print(success_result)
+
+    error_result = CollectorResult.error("An error occurred")
+    print(error_result)
+
+    unwrapped_value = success_result.unwrap()
+    print("Unwrapped:", unwrapped_value)
+
+    default_value = success_result.unwrap_or("Default value")
+    print("Unwrapped with default:", default_value)
+
+    print("Equality:", success_result == "Success value")
+    print("Equality:", success_result == CollectorResult.success("Success value"))

querent/collectors/collector_result.py

@@ -0,0 +1,54 @@
+import asyncio
+from pathlib import Path
+from typing import AsyncGenerator
+from querent.collectors.collector_base import Collector
+from querent.collectors.collector_factory import CollectorFactory
+from querent.collectors.collector_result import CollectorResult
+from querent.common.uri import Uri
+from querent.config.collector_config import CollectorBackend, FSCollectorConfig
+import aiofiles
+
+
+class FSCollector(Collector):
+    def __init__(self, config: FSCollectorConfig):
+        self.root_dir = Path(config.root_path)
+        self.chunk_size = config.chunk_size
+
+    async def connect(self):
+        # Add your setup logic here if needed
+        pass
+
+    async def disconnect(self):
+        # Add your cleanup logic here if needed
+        pass
+
+    async def poll(self) -> AsyncGenerator[CollectorResult, None]:
+        async for file_path in self.walk_files(self.root_dir):
+            async with aiofiles.open(file_path, "rb") as file:
+                async for chunk in self.read_chunks(file):
+                    yield CollectorResult({"file_path": file_path, "chunk": chunk})
+
+    async def read_chunks(self, file):
+        while True:
+            chunk = await file.read(self.chunk_size)
+            if not chunk:
+                break
+            yield chunk
+        await file.close()
+
+    async def walk_files(self, root: Path) -> AsyncGenerator[Path, None]:
+        for item in root.iterdir():
+            if item.is_file():
+                yield item
+
+
+class FSCollectorFactory(CollectorFactory):
+    def __init__(self):
+        pass
+
+    def backend(self) -> CollectorBackend:
+        return CollectorBackend.LocalFile
+
+    def resolve(self, uri: Uri) -> Collector:
+        config = FSCollectorConfig(root_path=uri.path)
+        return FSCollector(config)

querent/connectors/fs/__init__.py renamed to querent/collectors/fs/__init__.py

@@ -0,0 +1,53 @@
+from enum import Enum
+from typing import Optional
+from pydantic import BaseModel
+
+
+class CollectorBackend(str, Enum):
+    LocalFile = "localfile"
+    S3 = "s3"
+    Gcs = "gs"
+
+
+class CollectConfig(BaseModel):
+    backend: CollectorBackend
+
+    class Config:
+        use_enum_values = True
+
+
+class FSCollectorConfig(BaseModel):
+    root_path: str
+    chunk_size: int = 1024
+
+
+class S3CollectConfig(BaseModel):
+    bucket: str
+    region: str
+    access_key: str
+    secret_key: str
+
+
+class GcsCollectConfig(BaseModel):
+    bucket: str
+    region: str
+    access_key: str
+    secret_key: str
+
+
+class CollectConfigWrapper(BaseModel):
+    backend: CollectorBackend
+    config: Optional[BaseModel] = None
+
+    @classmethod
+    def from_collect_config(cls, collect_config: CollectConfig):
+        if collect_config.backend == CollectorBackend.LocalFile:
+            return cls(
+                backend=CollectorBackend.LocalFile, config=FSCollectorConfig()
+            )
+        elif collect_config.backend == CollectorBackend.S3:
+            return cls(backend=CollectorBackend.S3, config=S3CollectConfig())
+        elif collect_config.backend == CollectorBackend.Gcs:
+            return cls(backend=CollectorBackend.Gcs, config=GcsCollectConfig())
+        else:
+            raise ValueError(f"Unsupported collector backend: {collect_config.backend}")

querent/collectors/fs/fs_collector.py

@@ -154,3 +154,4 @@ html2text==2020.1.16
 duckduckgo-search==3.8.3 
 google-generativeai==0.1.0
 asyncio==3.4.3
+aiofiles