Change template to private attribute and simplify imports (#101)

---------

Co-authored-by: ian <ian@cinnamon.is>

knowledgehub/agents/base.py

@@ -1,8 +1,6 @@
 from typing import Optional, Union
 
-from theflow import Node, Param
-
-from kotaemon.base import BaseComponent
+from kotaemon.base import BaseComponent, Node, Param
 from kotaemon.llms import BaseLLM, PromptTemplate
 
 from .io import AgentOutput, AgentType

knowledgehub/agents/react/agent.py

@@ -2,11 +2,10 @@ import logging
 import re
 from typing import Optional
 
-from theflow import Param
-
 from kotaemon.agents.base import BaseAgent, BaseLLM
 from kotaemon.agents.io import AgentAction, AgentFinish, AgentOutput, AgentType
 from kotaemon.agents.tools import BaseTool
+from kotaemon.base import Param
 from kotaemon.llms import PromptTemplate
 
 FINAL_ANSWER_ACTION = "Final Answer:"

knowledgehub/agents/rewoo/agent.py

@@ -3,12 +3,11 @@ import re
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
-from theflow import Node, Param
-
 from kotaemon.agents.base import BaseAgent
 from kotaemon.agents.io import AgentOutput, AgentType, BaseScratchPad
 from kotaemon.agents.tools import BaseTool
 from kotaemon.agents.utils import get_plugin_response_content
+from kotaemon.base import Node, Param
 from kotaemon.indices.qa import CitationPipeline
 from kotaemon.llms import BaseLLM, PromptTemplate
 

knowledgehub/base/__init__.py

@@ -1,4 +1,4 @@
-from .component import BaseComponent
+from .component import BaseComponent, Node, Param, lazy
 from .schema import (
     AIMessage,
     BaseMessage,
@@ -22,4 +22,7 @@ __all__ = [
     "RetrievedDocument",
     "LLMInterface",
     "ExtractorOutput",
+    "Param",
+    "Node",
+    "lazy",
 ]

knowledgehub/base/component.py

@@ -1,6 +1,6 @@
 from abc import abstractmethod
 
-from theflow.base import Function
+from theflow import Function, Node, Param, lazy
 
 from kotaemon.base.schema import Document
 
@@ -35,3 +35,6 @@ class BaseComponent(Function):
     def run(self, *args, **kwargs) -> Document | list[Document] | None:
         """Run the component."""
         ...
+
+
+__all__ = ["BaseComponent", "Param", "Node", "lazy"]

knowledgehub/indices/ingests/files.py

@@ -1,9 +1,8 @@
 from pathlib import Path
 
 from llama_index.readers.base import BaseReader
-from theflow import Param
 
-from kotaemon.base import BaseComponent, Document
+from kotaemon.base import BaseComponent, Document, Param
 from kotaemon.indices.extractors import BaseDocParser
 from kotaemon.indices.splitters import BaseSplitter, TokenSplitter
 from kotaemon.loaders import (

knowledgehub/llms/branching.py

@@ -1,8 +1,7 @@
 from typing import List, Optional
 
-from theflow import Param
+from kotaemon.base import BaseComponent, Document, Param
 
-from ..base import BaseComponent, Document
 from .linear import GatedLinearPipeline
 
 

knowledgehub/llms/cot.py

@@ -1,9 +1,7 @@
 from copy import deepcopy
 from typing import Callable, List
 
-from theflow import Function, Node, Param
-
-from kotaemon.base import BaseComponent, Document
+from kotaemon.base import BaseComponent, Document, Node, Param
 
 from .chats import AzureChatOpenAI
 from .completions import LLM
@@ -74,7 +72,7 @@ class Thought(BaseComponent):
         )
     )
     llm: LLM = Node(AzureChatOpenAI, help="The LLM model to execute the input prompt")
-    post_process: Function = Node(
+    post_process: BaseComponent = Node(
         help=(
             "The function post-processor that post-processes LLM output prediction ."
             "It should take a string as input (this is the LLM output text) and return "
@@ -85,7 +83,7 @@ class Thought(BaseComponent):
     @Node.auto(depends_on="prompt")
     def prompt_template(self):
         """Automatically wrap around param prompt. Can ignore"""
-        return BasePromptComponent(self.prompt)
+        return BasePromptComponent(template=self.prompt)
 
     def run(self, **kwargs) -> Document:
         """Run the chain of thought"""

knowledgehub/llms/prompts/base.py

@@ -1,6 +1,7 @@
 from typing import Callable, Union
 
-from ...base import BaseComponent, Document
+from kotaemon.base import BaseComponent, Document
+
 from .template import PromptTemplate
 
 
@@ -16,6 +17,7 @@ class BasePromptComponent(BaseComponent):
 
     class Config:
         middleware_switches = {"theflow.middleware.CachingMiddleware": False}
+        allow_extra = True
 
     def __init__(self, template: Union[str, PromptTemplate], **kwargs):
         super().__init__()

knowledgehub/parsers/regex_extractor.py

@@ -3,10 +3,7 @@ from __future__ import annotations
 import re
 from typing import Callable
 
-from theflow import Param
-
-from kotaemon.base import BaseComponent, Document
-from kotaemon.base.schema import ExtractorOutput
+from kotaemon.base import BaseComponent, Document, ExtractorOutput, Param
 
 
 class RegexExtractor(BaseComponent):