Move Document and other interface into base/schema (#69)

knowledgehub/base/__init__.py

@@ -1,3 +1,4 @@
 from .component import BaseComponent
+from .schema import Document
 
-__all__ = ["BaseComponent"]
+__all__ = ["BaseComponent", "Document"]

knowledgehub/base/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Any, Optional, TypeVar
 
 from llama_index.bridge.pydantic import Field
@@ -72,3 +74,19 @@ class RetrievedDocument(Document):
 
     score: float = Field(default=0.0)
     retrieval_metadata: dict = Field(default={})
+
+
+class LLMInterface(Document):
+    candidates: list[str] = Field(default_factory=list)
+    completion_tokens: int = -1
+    total_tokens: int = -1
+    prompt_tokens: int = -1
+    logits: list[list[float]] = Field(default_factory=list)
+
+
+class ExtractorOutput(Document):
+    """
+    Represents the output of an extractor.
+    """
+
+    matches: list[str]

knowledgehub/chatbot/base.py

@@ -5,7 +5,7 @@ from langchain.schema.messages import AIMessage, SystemMessage
 from theflow import Param, SessionCompose
 
 from ..base import BaseComponent
-from ..llms.base import LLMInterface
+from ..base.schema import LLMInterface
 from ..llms.chats.base import BaseMessage, HumanMessage
 
 

knowledgehub/composite/branching.py

@@ -2,9 +2,8 @@ from typing import List, Optional
 
 from theflow import Param
 
-from kotaemon.base import BaseComponent
+from kotaemon.base import BaseComponent, Document
 from kotaemon.composite.linear import GatedLinearPipeline
-from kotaemon.documents.base import Document
 
 
 class SimpleBranchingPipeline(BaseComponent):

knowledgehub/composite/linear.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable, Optional, Union
 
 from kotaemon.base import BaseComponent
-from kotaemon.documents.base import Document, IO_Type
+from kotaemon.base.schema import Document, IO_Type
 from kotaemon.llms.chats.base import ChatLLM
 from kotaemon.llms.completions.base import LLM
 from kotaemon.prompt.base import BasePromptComponent

knowledgehub/docstores/base.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import List, Optional, Union
 
-from ..documents.base import Document
+from ..base import Document
 
 
 class BaseDocumentStore(ABC):

knowledgehub/docstores/in_memory.py

@@ -2,7 +2,7 @@ import json
 from pathlib import Path
 from typing import List, Optional, Union
 
-from ..documents.base import Document
+from ..base import Document
 from .base import BaseDocumentStore
 
 

knowledgehub/documents/__init__.py

@@ -1,3 +0,0 @@
-from .base import Document
-
-__all__ = ["Document"]

knowledgehub/embeddings/base.py

@@ -6,8 +6,7 @@ from typing import Type
 from langchain.schema.embeddings import Embeddings as LCEmbeddings
 from theflow import Param
 
-from ..base import BaseComponent
-from ..documents.base import Document
+from ..base import BaseComponent, Document
 
 
 class BaseEmbeddings(BaseComponent):

knowledgehub/llms/base.py

@@ -1,13 +0,0 @@
-from typing import List
-
-from pydantic import Field
-
-from kotaemon.documents.base import Document
-
-
-class LLMInterface(Document):
-    candidates: List[str] = Field(default_factory=list)
-    completion_tokens: int = -1
-    total_tokens: int = -1
-    prompt_tokens: int = -1
-    logits: List[List[float]] = Field(default_factory=list)

knowledgehub/llms/chats/base.py

@@ -8,7 +8,7 @@ from langchain.schema.messages import BaseMessage, HumanMessage
 from theflow.base import Param
 
 from ...base import BaseComponent
-from ..base import LLMInterface
+from ...base.schema import LLMInterface
 
 logger = logging.getLogger(__name__)
 

knowledgehub/llms/completions/base.py

@@ -5,7 +5,7 @@ from langchain.llms.base import BaseLLM
 from theflow.base import Param
 
 from ...base import BaseComponent
-from ..base import LLMInterface
+from ...base.schema import LLMInterface
 
 logger = logging.getLogger(__name__)
 

knowledgehub/loaders/base.py

@@ -4,8 +4,7 @@ from typing import Any, List, Type, Union
 from llama_index import SimpleDirectoryReader, download_loader
 from llama_index.readers.base import BaseReader
 
-from ..base import BaseComponent
-from ..documents.base import Document
+from ..base import BaseComponent, Document
 
 
 class AutoReader(BaseComponent):

knowledgehub/loaders/excel_loader.py

@@ -8,7 +8,7 @@ from typing import Any, List, Optional, Union
 
 from llama_index.readers.base import BaseReader
 
-from kotaemon.documents import Document
+from kotaemon.base import Document
 
 
 class PandasExcelReader(BaseReader):

knowledgehub/loaders/mathpix_loader.py

@@ -8,7 +8,7 @@ import requests
 from langchain.utils import get_from_dict_or_env
 from llama_index.readers.base import BaseReader
 
-from kotaemon.documents import Document
+from kotaemon.base import Document
 
 from .utils.table import parse_markdown_text_to_tables, strip_special_chars_markdown
 

knowledgehub/loaders/ocr_loader.py

@@ -5,7 +5,7 @@ from uuid import uuid4
 import requests
 from llama_index.readers.base import BaseReader
 
-from kotaemon.documents import Document
+from kotaemon.base import Document
 
 from .utils.pdf_ocr import parse_ocr_output, read_pdf_unstructured
 from .utils.table import strip_special_chars_markdown

knowledgehub/parsers/splitter.py

@@ -4,9 +4,7 @@ from llama_index.node_parser import SimpleNodeParser as LISimpleNodeParser
 from llama_index.node_parser.interface import NodeParser
 from llama_index.text_splitter import TokenTextSplitter
 
-from kotaemon.base import BaseComponent
-
-from ..documents.base import Document
+from ..base import BaseComponent, Document
 
 __all__ = ["TokenTextSplitter"]
 

knowledgehub/pipelines/indexing.py

@@ -5,9 +5,8 @@ from pathlib import Path
 
 from theflow import Node, Param
 
-from ..base import BaseComponent
+from ..base import BaseComponent, Document
 from ..docstores import BaseDocumentStore
-from ..documents.base import Document
 from ..embeddings import BaseEmbeddings
 from ..vectorstores import BaseVectorStore
 

knowledgehub/pipelines/qa.py

@@ -6,8 +6,8 @@ from theflow import Node
 from theflow.utils.modules import ObjectInitDeclaration as _
 
 from kotaemon.base import BaseComponent
+from kotaemon.base.schema import RetrievedDocument
 from kotaemon.docstores import InMemoryDocumentStore
-from kotaemon.documents.base import RetrievedDocument
 from kotaemon.embeddings import AzureOpenAIEmbeddings
 from kotaemon.llms.chats.openai import AzureChatOpenAI
 from kotaemon.pipelines.agents import BaseAgent

knowledgehub/pipelines/retrieving.py

@@ -5,8 +5,8 @@ from pathlib import Path
 from theflow import Node, Param
 
 from ..base import BaseComponent
+from ..base.schema import Document, RetrievedDocument
 from ..docstores import BaseDocumentStore
-from ..documents.base import Document, RetrievedDocument
 from ..embeddings import BaseEmbeddings
 from ..vectorstores import BaseVectorStore
 

knowledgehub/pipelines/tools/wikipedia.py

@@ -2,7 +2,7 @@ from typing import Any, AnyStr, Optional, Type, Union
 
 from pydantic import BaseModel, Field
 
-from kotaemon.documents.base import Document
+from kotaemon.base import Document
 
 from .base import BaseTool
 

knowledgehub/post_processing/extractor.py

@@ -5,16 +5,8 @@ from typing import Callable
 
 from theflow import Param
 
-from kotaemon.base import BaseComponent
-from kotaemon.documents.base import Document
-
-
-class ExtractorOutput(Document):
-    """
-    Represents the output of an extractor.
-    """
-
-    matches: list[str]
+from kotaemon.base import BaseComponent, Document
+from kotaemon.base.schema import ExtractorOutput
 
 
 class RegexExtractor(BaseComponent):

knowledgehub/prompt/base.py

@@ -1,7 +1,6 @@
 from typing import Callable, Union
 
-from kotaemon.base import BaseComponent
-from kotaemon.documents.base import Document
+from kotaemon.base import BaseComponent, Document
 from kotaemon.prompt.template import PromptTemplate
 
 

knowledgehub/vectorstores/base.py

@@ -6,7 +6,7 @@ from llama_index.vector_stores.types import BasePydanticVectorStore
 from llama_index.vector_stores.types import VectorStore as LIVectorStore
 from llama_index.vector_stores.types import VectorStoreQuery
 
-from ..documents.base import Document
+from ..base import Document
 
 
 class BaseVectorStore(ABC):