Skip to content

Commit

Permalink
feat(agent):merge new code for branch main
Browse files Browse the repository at this point in the history
  • Loading branch information
yhjun1026 committed Jun 13, 2024
2 parents ca04294 + 58d0878 commit 3adb7e9
Show file tree
Hide file tree
Showing 96 changed files with 1,105 additions and 572 deletions.
4 changes: 2 additions & 2 deletions .env.template
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,15 @@
#** LLM MODELS **#
#*******************************************************************#
# LLM_MODEL, see dbgpt/configs/model_config.LLM_MODEL_CONFIG
LLM_MODEL=vicuna-13b-v1.5
LLM_MODEL=glm-4-9b-chat
## LLM model path, by default, DB-GPT will read the model path from LLM_MODEL_CONFIG based on the LLM_MODEL.
## Of course you can specify your model path according to LLM_MODEL_PATH
## In DB-GPT, the priority from high to low to read model path:
## 1. environment variable with key: {LLM_MODEL}_MODEL_PATH (Avoid multi-model conflicts)
## 2. environment variable with key: MODEL_PATH
## 3. environment variable with key: LLM_MODEL_PATH
## 4. the config in dbgpt/configs/model_config.LLM_MODEL_CONFIG
# LLM_MODEL_PATH=/app/models/vicuna-13b-v1.5
# LLM_MODEL_PATH=/app/models/glm-4-9b-chat
# LLM_PROMPT_TEMPLATE=vicuna_v1.1
MODEL_SERVER=http://127.0.0.1:8000
LIMIT_MODEL_CONCURRENCY=5
Expand Down
2 changes: 1 addition & 1 deletion dbgpt/_private/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def __init__(self) -> None:
self.CHAT_HISTORY_STORE_TYPE = os.getenv("CHAT_HISTORY_STORE_TYPE", "db")

### LLM Model Service Configuration
self.LLM_MODEL = os.getenv("LLM_MODEL", "vicuna-13b-v1.5")
self.LLM_MODEL = os.getenv("LLM_MODEL", "glm-4-9b-chat")
self.LLM_MODEL_PATH = os.getenv("LLM_MODEL_PATH")

### Proxy llm backend, this configuration is only valid when "LLM_MODEL=proxyllm"
Expand Down
29 changes: 14 additions & 15 deletions dbgpt/agent/core/memory/hybrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,10 @@
import os.path
from concurrent.futures import Executor, ThreadPoolExecutor
from datetime import datetime
from typing import TYPE_CHECKING, Generic, List, Optional, Tuple, Type
from typing import Generic, List, Optional, Tuple, Type

from dbgpt.core import Embeddings, LLMClient
from dbgpt.storage.vector_store.base import VectorStoreBase
from dbgpt.util.annotations import immutable, mutable

from .base import (
Expand All @@ -26,9 +27,6 @@
from .long_term import LongTermMemory
from .short_term import EnhancedShortTermMemory

if TYPE_CHECKING:
from dbgpt.storage.vector_store.connector import VectorStoreConnector


class HybridMemory(Memory, Generic[T]):
"""Hybrid memory for the agent."""
Expand Down Expand Up @@ -81,8 +79,10 @@ def from_chroma(
):
"""Create a hybrid memory from Chroma vector store."""
from dbgpt.configs.model_config import DATA_DIR
from dbgpt.storage.vector_store.chroma_store import ChromaVectorConfig
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt.storage.vector_store.chroma_store import (
ChromaStore,
ChromaVectorConfig,
)

if not embeddings:
from dbgpt.rag.embedding import DefaultEmbeddingFactory
Expand All @@ -91,16 +91,15 @@ def from_chroma(

vstore_path = vstore_path or os.path.join(DATA_DIR, "agent_memory")

vector_store_connector = VectorStoreConnector.from_default(
vector_store_type="Chroma",
embedding_fn=embeddings,
vector_store_config=ChromaVectorConfig(
vector_store = ChromaStore(
ChromaVectorConfig(
name=vstore_name,
persist_path=vstore_path,
),
embedding_fn=embeddings,
)
)
return cls.from_vstore(
vector_store_connector=vector_store_connector,
vector_store=vector_store,
embeddings=embeddings,
executor=executor,
now=now,
Expand All @@ -113,7 +112,7 @@ def from_chroma(
@classmethod
def from_vstore(
cls,
vector_store_connector: "VectorStoreConnector",
vector_store: "VectorStoreBase",
embeddings: Optional[Embeddings] = None,
executor: Optional[Executor] = None,
now: Optional[datetime] = None,
Expand All @@ -124,7 +123,7 @@ def from_vstore(
):
"""Create a hybrid memory from vector store."""
if not embeddings:
embeddings = vector_store_connector.current_embeddings
raise ValueError("embeddings is required.")
if not executor:
executor = ThreadPoolExecutor()
if not now:
Expand All @@ -139,7 +138,7 @@ def from_vstore(
if not long_term_memory:
long_term_memory = LongTermMemory(
executor,
vector_store_connector,
vector_store,
now=now,
)
return cls(now, sensory_memory, short_term_memory, long_term_memory, **kwargs)
Expand Down
10 changes: 5 additions & 5 deletions dbgpt/agent/core/memory/long_term.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

from dbgpt.core import Chunk
from dbgpt.rag.retriever.time_weighted import TimeWeightedEmbeddingRetriever
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt.storage.vector_store.base import VectorStoreBase
from dbgpt.storage.vector_store.filters import MetadataFilters
from dbgpt.util.annotations import immutable, mutable
from dbgpt.util.executor_utils import blocking_func_to_async
Expand Down Expand Up @@ -70,7 +70,7 @@ class LongTermMemory(Memory, Generic[T]):
def __init__(
self,
executor: Executor,
vector_store_connector: VectorStoreConnector,
vector_store: VectorStoreBase,
now: Optional[datetime] = None,
reflection_threshold: Optional[float] = None,
):
Expand All @@ -81,9 +81,9 @@ def __init__(
self.forgetting: bool = False
self.reflection_threshold: Optional[float] = reflection_threshold
self.aggregate_importance: float = 0.0
self._vector_store_connector = vector_store_connector
self._vector_store = vector_store
self.memory_retriever = LongTermRetriever(
now=self.now, vector_store_connector=vector_store_connector
now=self.now, index_store=vector_store
)

@immutable
Expand All @@ -97,7 +97,7 @@ def structure_clone(
m: LongTermMemory[T] = LongTermMemory(
now=now,
executor=self.executor,
vector_store_connector=self._vector_store_connector.new_connector(new_name),
vector_store=self._vector_store,
reflection_threshold=self.reflection_threshold,
)
m._copy_from(self)
Expand Down
2 changes: 2 additions & 0 deletions dbgpt/agent/expand/resources/search_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ def baidu_search(
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:112.0) "
"Gecko/20100101 Firefox/112.0"
}
if num_results < 8:
num_results = 8
url = f"https://www.baidu.com/s?wd={query}&rn={num_results}"
response = requests.get(url, headers=headers)
response.encoding = "utf-8"
Expand Down
4 changes: 2 additions & 2 deletions dbgpt/app/knowledge/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@
from dbgpt.rag.knowledge.factory import KnowledgeFactory
from dbgpt.rag.retriever.embedding import EmbeddingRetriever
from dbgpt.serve.rag.api.schemas import KnowledgeSyncRequest
from dbgpt.serve.rag.connector import VectorStoreConnector
from dbgpt.serve.rag.service.service import Service
from dbgpt.storage.vector_store.base import VectorStoreConfig
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt.util.tracer import SpanType, root_tracer

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -310,7 +310,7 @@ def similar_query(space_name: str, query_request: KnowledgeQueryRequest):
vector_store_config=config,
)
retriever = EmbeddingRetriever(
top_k=query_request.top_k, vector_store_connector=vector_store_connector
top_k=query_request.top_k, index_store=vector_store_connector.index_client
)
chunks = retriever.retrieve(query_request.query)
res = [
Expand Down
6 changes: 2 additions & 4 deletions dbgpt/app/knowledge/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
ChunkQueryRequest,
DocumentQueryRequest,
DocumentSummaryRequest,
DocumentSyncRequest,
KnowledgeDocumentRequest,
KnowledgeSpaceRequest,
SpaceArgumentRequest,
Expand All @@ -24,19 +23,18 @@
)
from dbgpt.component import ComponentType
from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG
from dbgpt.core import Chunk, LLMClient
from dbgpt.core import LLMClient
from dbgpt.model import DefaultLLMClient
from dbgpt.model.cluster import WorkerManagerFactory
from dbgpt.rag.assembler.embedding import EmbeddingAssembler
from dbgpt.rag.assembler.summary import SummaryAssembler
from dbgpt.rag.chunk_manager import ChunkParameters
from dbgpt.rag.embedding.embedding_factory import EmbeddingFactory
from dbgpt.rag.knowledge.base import KnowledgeType
from dbgpt.rag.knowledge.factory import KnowledgeFactory
from dbgpt.serve.rag.connector import VectorStoreConnector
from dbgpt.serve.rag.models.models import KnowledgeSpaceDao, KnowledgeSpaceEntity
from dbgpt.serve.rag.service.service import SyncStatus
from dbgpt.storage.vector_store.base import VectorStoreConfig
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt.util.executor_utils import ExecutorFactory, blocking_func_to_async
from dbgpt.util.tracer import root_tracer, trace

Expand Down
4 changes: 2 additions & 2 deletions dbgpt/app/scene/chat_knowledge/v1/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def __init__(self, chat_param: Dict):
"embedding_factory", EmbeddingFactory
)
from dbgpt.rag.retriever.embedding import EmbeddingRetriever
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt.serve.rag.connector import VectorStoreConnector

embedding_fn = embedding_factory.create(
model_name=EMBEDDING_MODEL_CONFIG[CFG.EMBEDDING_MODEL]
Expand Down Expand Up @@ -116,7 +116,7 @@ def __init__(self, chat_param: Dict):
retriever_top_k = max(CFG.RERANK_TOP_K, 20)
self.embedding_retriever = EmbeddingRetriever(
top_k=retriever_top_k,
vector_store_connector=vector_store_connector,
index_store=vector_store_connector.index_client,
query_rewrite=query_rewrite,
rerank=reranker,
)
Expand Down
2 changes: 1 addition & 1 deletion dbgpt/app/static/404.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/404/index.html

Large diffs are not rendered by default.

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Large diffs are not rendered by default.

This file was deleted.

2 changes: 1 addition & 1 deletion dbgpt/app/static/agent/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/app/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/chat/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/database/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/flow/canvas/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/flow/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/knowledge/chunk/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/knowledge/graph/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/knowledge/index.html

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion dbgpt/app/static/models/index.html

Large diffs are not rendered by default.

Binary file modified dbgpt/app/static/models/knowledge-default.jpg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added dbgpt/app/static/models/knowledge-full-text.jpg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
2 changes: 1 addition & 1 deletion dbgpt/app/static/prompt/index.html

Large diffs are not rendered by default.

41 changes: 10 additions & 31 deletions dbgpt/rag/assembler/db_schema.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
"""DBSchemaAssembler."""
from typing import Any, List, Optional

from dbgpt.core import Chunk, Embeddings
from dbgpt.core import Chunk
from dbgpt.datasource.base import BaseConnector
from dbgpt.storage.vector_store.connector import VectorStoreConnector

from ..assembler.base import BaseAssembler
from ..chunk_manager import ChunkParameters
from ..embedding.embedding_factory import DefaultEmbeddingFactory
from ..index.base import IndexStoreBase
from ..knowledge.datasource import DatasourceKnowledge
from ..retriever.db_schema import DBSchemaRetriever

Expand Down Expand Up @@ -36,36 +35,22 @@ class DBSchemaAssembler(BaseAssembler):
def __init__(
self,
connector: BaseConnector,
vector_store_connector: VectorStoreConnector,
index_store: IndexStoreBase,
chunk_parameters: Optional[ChunkParameters] = None,
embedding_model: Optional[str] = None,
embeddings: Optional[Embeddings] = None,
**kwargs: Any,
) -> None:
"""Initialize with Embedding Assembler arguments.
Args:
connector: (BaseConnector) BaseConnector connection.
vector_store_connector: (VectorStoreConnector) VectorStoreConnector to use.
index_store: (IndexStoreBase) IndexStoreBase to use.
chunk_manager: (Optional[ChunkManager]) ChunkManager to use for chunking.
embedding_model: (Optional[str]) Embedding model to use.
embeddings: (Optional[Embeddings]) Embeddings to use.
"""
knowledge = DatasourceKnowledge(connector)
self._connector = connector
self._vector_store_connector = vector_store_connector

self._embedding_model = embedding_model
if self._embedding_model and not embeddings:
embeddings = DefaultEmbeddingFactory(
default_model_name=self._embedding_model
).create(self._embedding_model)

if (
embeddings
and self._vector_store_connector.vector_store_config.embedding_fn is None
):
self._vector_store_connector.vector_store_config.embedding_fn = embeddings
self._index_store = index_store

super().__init__(
knowledge=knowledge,
Expand All @@ -77,29 +62,23 @@ def __init__(
def load_from_connection(
cls,
connector: BaseConnector,
vector_store_connector: VectorStoreConnector,
index_store: IndexStoreBase,
chunk_parameters: Optional[ChunkParameters] = None,
embedding_model: Optional[str] = None,
embeddings: Optional[Embeddings] = None,
) -> "DBSchemaAssembler":
"""Load document embedding into vector store from path.
Args:
connector: (BaseConnector) BaseConnector connection.
vector_store_connector: (VectorStoreConnector) VectorStoreConnector to use.
index_store: (IndexStoreBase) IndexStoreBase to use.
chunk_parameters: (Optional[ChunkParameters]) ChunkManager to use for
chunking.
embedding_model: (Optional[str]) Embedding model to use.
embeddings: (Optional[Embeddings]) Embeddings to use.
Returns:
DBSchemaAssembler
"""
return cls(
connector=connector,
vector_store_connector=vector_store_connector,
embedding_model=embedding_model,
index_store=index_store,
chunk_parameters=chunk_parameters,
embeddings=embeddings,
)

def get_chunks(self) -> List[Chunk]:
Expand All @@ -112,7 +91,7 @@ def persist(self) -> List[str]:
Returns:
List[str]: List of chunk ids.
"""
return self._vector_store_connector.load_document(self._chunks)
return self._index_store.load_document(self._chunks)

def _extract_info(self, chunks) -> List[Chunk]:
"""Extract info from chunks."""
Expand All @@ -131,5 +110,5 @@ def as_retriever(self, top_k: int = 4, **kwargs) -> DBSchemaRetriever:
top_k=top_k,
connector=self._connector,
is_embeddings=True,
vector_store_connector=self._vector_store_connector,
index_store=self._index_store,
)
Loading

0 comments on commit 3adb7e9

Please sign in to comment.