Merge pull request #54 from CambioML/refactor

Refactor fix

example/flex/flex_demo.py

@@ -2,13 +2,17 @@
 
 import os
 import argparse
+from dotenv import load_dotenv
 from pykoi import Application
 from pykoi.chat import QuestionAnswerDatabase
 from pykoi.retrieval import RetrievalFactory
 from pykoi.retrieval import VectorDbFactory
 from pykoi.component import Chatbot, Dashboard, RetrievalQA
 
 
+load_dotenv()
+
+
 def main(**kargs):
     os.environ["DOC_PATH"] = os.path.join(os.getcwd(), "temp/docs")
     os.environ["VECTORDB_PATH"] = os.path.join(os.getcwd(), "temp/vectordb")

pykoi/application.py

@@ -14,7 +14,6 @@
 from pydantic import BaseModel
 from pyngrok import ngrok
 from starlette.middleware.cors import CORSMiddleware
-from pykoi.component.base import Dropdown
 from pykoi.interactives.chatbot import Chatbot
 from pykoi.telemetry.telemetry import Telemetry
 from pykoi.telemetry.events import AppStartEvent, AppStopEvent
@@ -157,6 +156,9 @@ def add_component(self, component: Any):
         Args:
             component (Any): The component to be added.
         """
+        # TODO (Jojo): remove component Dropdown check
+        from pykoi.component.base import Dropdown
+
         if component.data_source:
             self.data_sources[component.id] = component.data_source
             # set data_endpoint if it's a Dropdown component

pykoi/retrieval/llm/embedding_factory.py

@@ -27,8 +27,10 @@ def create_embedding(model_source: Union[str, ModelSource], **kwargs) -> Embeddi
         try:
             model_source = ModelSource(model_source)
             if model_source == ModelSource.OPENAI:
+                from langchain.embeddings import OpenAIEmbeddings
                 return OpenAIEmbeddings()
             elif model_source == ModelSource.HUGGINGFACE:
+                from langchain.embeddings import HuggingFaceEmbeddings
                 return HuggingFaceEmbeddings(
                     model_name=kwargs.get("model_name"),
                 )

pykoi/retrieval/llm/openai.py

@@ -5,12 +5,10 @@
 
 from langchain.chains import RetrievalQA
 from langchain.llms import OpenAI
-from dotenv import load_dotenv
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
 
-load_dotenv()
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 MIN_DOCS = 2

pykoi/retrieval/llm/retrieval_factory.py

@@ -2,8 +2,6 @@
 from typing import Union
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
-from pykoi.retrieval.llm.openai import OpenAIModel
-from pykoi.retrieval.llm.huggingface import HuggingFaceModel
 from pykoi.retrieval.llm.constants import ModelSource
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
 
@@ -29,8 +27,10 @@ def create(
         try:
             model_source = ModelSource(model_source)
             if model_source == ModelSource.OPENAI:
+                from pykoi.retrieval.llm.openai import OpenAIModel
                 return OpenAIModel(vector_db)
             if model_source == ModelSource.HUGGINGFACE:
+                from pykoi.retrieval.llm.huggingface import HuggingFaceModel
                 return HuggingFaceModel(vector_db, **kwargs)
         except Exception as ex:
             raise Exception(f"Unknown model: {model_source}") from ex

pyproject.toml

@@ -21,6 +21,7 @@ bcrypt = "4.0.1"
 posthog = "3.0.1"
 pynvml = "11.5.0"
 pandas = "2.0.3"
+python-dotenv = "1.0.0"
 
 transformers = { version = "4.31.0", optional = true }
 einops = { version = "0.6.1", optional = true }
@@ -47,7 +48,8 @@ huggingface = [
     "transformers",
     "einops",
     "accelerate",
-    "bitsandbytes"
+    "bitsandbytes",
+    "sentence-transformers"
 ]
 rag = [
     "langchain", 
@@ -57,8 +59,7 @@ rag = [
     "pdfminer-six", 
     "docx2txt", 
     "python-multipart", 
-    "tiktoken", 
-    "sentence-transformers"
+    "tiktoken"
 ]
 rlhf = [
     "transformers",