fix: Fix examples/awel get model_name from model_config (#1112)

Co-authored-by: xiuzhu <[email protected]>
eosphoros-ai · Jan 24, 2024 · 8f18478 · 8f18478
1 parent f8c0064
commit 8f18478
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/examples/awel/simple_dbschema_retriever_example.py b/examples/awel/simple_dbschema_retriever_example.py
@@ -3,7 +3,8 @@
 
 from pydantic import BaseModel, Field
 
-from dbgpt.configs.model_config import MODEL_PATH, PILOT_PATH
+from dbgpt._private.config import Config
+from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG, MODEL_PATH, PILOT_PATH
 from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
 from dbgpt.datasource.rdbms.conn_sqlite import SQLiteTempConnect
 from dbgpt.rag.chunk import Chunk
@@ -38,17 +39,19 @@
 """
 
 
+CFG = Config()
+
+
 def _create_vector_connector():
     """Create vector connector."""
-    model_name = os.getenv("EMBEDDING_MODEL", "text2vec")
     return VectorStoreConnector.from_default(
         "Chroma",
         vector_store_config=ChromaVectorConfig(
             name="vector_name",
             persist_path=os.path.join(PILOT_PATH, "data"),
         ),
         embedding_fn=DefaultEmbeddingFactory(
-            default_model_name=os.path.join(MODEL_PATH, model_name),
+            default_model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL],
         ).create(),
     )
 

diff --git a/examples/awel/simple_rag_embedding_example.py b/examples/awel/simple_rag_embedding_example.py
@@ -3,7 +3,8 @@
 
 from pydantic import BaseModel, Field
 
-from dbgpt.configs.model_config import MODEL_PATH, PILOT_PATH
+from dbgpt._private.config import Config
+from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG, MODEL_PATH, PILOT_PATH
 from dbgpt.core.awel import DAG, HttpTrigger, MapOperator
 from dbgpt.rag.embedding.embedding_factory import DefaultEmbeddingFactory
 from dbgpt.rag.knowledge.base import KnowledgeType
@@ -25,18 +26,19 @@
             }'
 """
 
+CFG = Config()
+
 
 def _create_vector_connector() -> VectorStoreConnector:
     """Create vector connector."""
-    model_name = os.getenv("EMBEDDING_MODEL", "text2vec")
     return VectorStoreConnector.from_default(
         "Chroma",
         vector_store_config=ChromaVectorConfig(
             name="vector_name",
             persist_path=os.path.join(PILOT_PATH, "data"),
         ),
         embedding_fn=DefaultEmbeddingFactory(
-            default_model_name=os.path.join(MODEL_PATH, model_name),
+            default_model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL],
         ).create(),
     )
 

diff --git a/examples/awel/simple_rag_retriever_example.py b/examples/awel/simple_rag_retriever_example.py
@@ -4,7 +4,8 @@
 
 from pydantic import BaseModel, Field
 
-from dbgpt.configs.model_config import MODEL_PATH, PILOT_PATH
+from dbgpt._private.config import Config
+from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG, MODEL_PATH, PILOT_PATH
 from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
 from dbgpt.model.proxy import OpenAILLMClient
 from dbgpt.rag.chunk import Chunk
@@ -43,6 +44,8 @@
             }'
 """
 
+CFG = Config()
+
 
 class TriggerReqBody(BaseModel):
     query: str = Field(..., description="User query")
@@ -83,7 +86,7 @@ def _create_vector_connector():
             persist_path=os.path.join(PILOT_PATH, "data"),
         ),
         embedding_fn=DefaultEmbeddingFactory(
-            default_model_name=os.path.join(MODEL_PATH, model_name),
+            default_model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL],
         ).create(),
     )