from __future__ import annotations
import json
import logging
from hashlib import sha1
from threading import Thread
from typing import Any, Dict, Iterable, List, Optional, Tuple
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseSettings
from langchain_core.vectorstores import VectorStore
logger = logging.getLogger()
DEBUG = False
[docs]def has_mul_sub_str(s: str, *args: Any) -> bool:
"""检查字符串是否具有多个子字符串。
参数:
s:要检查的字符串
*args:要在字符串中检查的子字符串
返回:
bool:如果所有子字符串都存在于字符串中,则为True,否则为False
"""
for a in args:
if a not in s:
return False
return True
[docs]def debug_output(s: Any) -> None:
"""如果DEBUG为True,则打印调试消息。
参数:
s:要打印的消息
"""
if DEBUG:
print(s) # noqa: T201
[docs]def get_named_result(connection: Any, query: str) -> List[dict[str, Any]]:
"""从查询中获取一个命名结果。
参数:
connection:数据库连接
query:要执行的查询
返回:
List[dict[str, Any]]:查询的结果
"""
cursor = connection.cursor()
cursor.execute(query)
columns = cursor.description
result = []
for value in cursor.fetchall():
r = {}
for idx, datum in enumerate(value):
k = columns[idx][0]
r[k] = datum
result.append(r)
debug_output(result)
cursor.close()
return result
[docs]class StarRocksSettings(BaseSettings):
"""StarRocks客户端配置。
属性:
StarRocks_host (str) : 用于连接到MyScale后端的URL。
默认为'localhost'。
StarRocks_port (int) : 用于通过HTTP连接的URL端口。默认为8443。
username (str) : 登录的用户名。默认为None。
password (str) : 登录的密码。默认为None。
database (str) : 要查找表的数据库名称。默认为'default'。
table (str) : 要操作的表名。
默认为'vector_table'。
column_map (Dict) : 列类型映射,将列名投影到langchain语义上。
必须具有键:`text`、`id`、`vector`,
必须与列数相同。例如:
.. code-block:: python
{
'id': 'text_id',
'embedding': 'text_embedding',
'document': 'text_plain',
'metadata': 'metadata_dictionary_in_json',
}
默认为身份映射。"""
host: str = "localhost"
port: int = 9030
username: str = "root"
password: str = ""
column_map: Dict[str, str] = {
"id": "id",
"document": "document",
"embedding": "embedding",
"metadata": "metadata",
}
database: str = "default"
table: str = "langchain"
def __getitem__(self, item: str) -> Any:
return getattr(self, item)
class Config:
env_file = ".env"
env_prefix = "starrocks_"
env_file_encoding = "utf-8"
[docs]class StarRocks(VectorStore):
"""`StarRocks` 向量存储。
您需要一个 `pymysql` python 包,以及一个有效的账户
来连接到 StarRocks。
目前 StarRocks 只实现了 `cosine_similarity` 函数
用于计算两个向量之间的距离。目前里面没有向量,
因此我们必须迭代所有向量并计算空间距离。
更多信息,请访问
[StarRocks 官方网站](https://www.starrocks.io/)
[StarRocks github](https://github.com/StarRocks/starrocks)"""
[docs] def __init__(
self,
embedding: Embeddings,
config: Optional[StarRocksSettings] = None,
**kwargs: Any,
) -> None:
"""StarRocks包装器到LangChain
embedding_function (嵌入):
config (StarRocks设置): StarRocks客户端的配置
"""
try:
import pymysql # type: ignore[import]
except ImportError:
raise ImportError(
"Could not import pymysql python package. "
"Please install it with `pip install pymysql`."
)
try:
from tqdm import tqdm
self.pgbar = tqdm
except ImportError:
# Just in case if tqdm is not installed
self.pgbar = lambda x, **kwargs: x
super().__init__()
if config is not None:
self.config = config
else:
self.config = StarRocksSettings()
assert self.config
assert self.config.host and self.config.port
assert self.config.column_map and self.config.database and self.config.table
for k in ["id", "embedding", "document", "metadata"]:
assert k in self.config.column_map
# initialize the schema
dim = len(embedding.embed_query("test"))
self.schema = f"""\
CREATE TABLE IF NOT EXISTS {self.config.database}.{self.config.table}(
{self.config.column_map['id']} string,
{self.config.column_map['document']} string,
{self.config.column_map['embedding']} array<float>,
{self.config.column_map['metadata']} string
) ENGINE = OLAP PRIMARY KEY(id) DISTRIBUTED BY HASH(id) \
PROPERTIES ("replication_num" = "1")\
"""
self.dim = dim
self.BS = "\\"
self.must_escape = ("\\", "'")
self.embedding_function = embedding
self.dist_order = "DESC"
debug_output(self.config)
# Create a connection to StarRocks
self.connection = pymysql.connect(
host=self.config.host,
port=self.config.port,
user=self.config.username,
password=self.config.password,
database=self.config.database,
**kwargs,
)
debug_output(self.schema)
get_named_result(self.connection, self.schema)
[docs] def escape_str(self, value: str) -> str:
return "".join(f"{self.BS}{c}" if c in self.must_escape else c for c in value)
@property
def embeddings(self) -> Embeddings:
return self.embedding_function
def _build_insert_sql(self, transac: Iterable, column_names: Iterable[str]) -> str:
ks = ",".join(column_names)
embed_tuple_index = tuple(column_names).index(
self.config.column_map["embedding"]
)
_data = []
for n in transac:
n = ",".join(
[
(
f"'{self.escape_str(str(_n))}'"
if idx != embed_tuple_index
else f"array<float>{str(_n)}"
)
for (idx, _n) in enumerate(n)
]
)
_data.append(f"({n})")
i_str = f"""
INSERT INTO
{self.config.database}.{self.config.table}({ks})
VALUES
{','.join(_data)}
"""
return i_str
def _insert(self, transac: Iterable, column_names: Iterable[str]) -> None:
_insert_query = self._build_insert_sql(transac, column_names)
debug_output(_insert_query)
get_named_result(self.connection, _insert_query)
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
batch_size: int = 32,
ids: Optional[Iterable[str]] = None,
**kwargs: Any,
) -> List[str]:
"""通过嵌入插入更多文本,并添加到VectorStore中。
参数:
texts:要添加到VectorStore中的字符串的可迭代对象。
ids:要与文本关联的可选id列表。
batch_size:插入的批量大小。
metadata:要插入的可选列数据。
返回:
将文本添加到VectorStore中的id列表。
"""
# Embed and create the documents
ids = ids or [sha1(t.encode("utf-8")).hexdigest() for t in texts]
colmap_ = self.config.column_map
transac = []
column_names = {
colmap_["id"]: ids,
colmap_["document"]: texts,
colmap_["embedding"]: self.embedding_function.embed_documents(list(texts)),
}
metadatas = metadatas or [{} for _ in texts]
column_names[colmap_["metadata"]] = map(json.dumps, metadatas)
assert len(set(colmap_) - set(column_names)) >= 0
keys, values = zip(*column_names.items())
try:
t = None
for v in self.pgbar(
zip(*values), desc="Inserting data...", total=len(metadatas)
):
assert (
len(v[keys.index(self.config.column_map["embedding"])]) == self.dim
)
transac.append(v)
if len(transac) == batch_size:
if t:
t.join()
t = Thread(target=self._insert, args=[transac, keys])
t.start()
transac = []
if len(transac) > 0:
if t:
t.join()
self._insert(transac, keys)
return [i for i in ids]
except Exception as e:
logger.error(f"\033[91m\033[1m{type(e)}\033[0m \033[95m{str(e)}\033[0m")
return []
[docs] @classmethod
def from_texts(
cls,
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[Dict[Any, Any]]] = None,
config: Optional[StarRocksSettings] = None,
text_ids: Optional[Iterable[str]] = None,
batch_size: int = 32,
**kwargs: Any,
) -> StarRocks:
"""使用现有文本创建StarRocks包装器
参数:
embedding_function (Embeddings): 用于提取文本嵌入的函数
texts (Iterable[str]): 要添加的字符串列表或元组
config (StarRocksSettings, Optional): StarRocks配置
text_ids (Optional[Iterable], optional): 文本的ID。默认为None。
batch_size (int, optional): 传输数据到StarRocks时的批处理大小。默认为32。
metadata (List[dict], optional): 文本的元数据。默认为None。
返回:
StarRocks索引
"""
ctx = cls(embedding, config, **kwargs)
ctx.add_texts(texts, ids=text_ids, batch_size=batch_size, metadatas=metadatas)
return ctx
def __repr__(self) -> str:
"""为StarRocks Vector Store提供文本表示,打印后端、用户名和模式。可通过`str(StarRocks())`轻松使用。
返回:
repr: 显示连接信息和数据模式的字符串
"""
_repr = f"\033[92m\033[1m{self.config.database}.{self.config.table} @ "
_repr += f"{self.config.host}:{self.config.port}\033[0m\n\n"
_repr += f"\033[1musername: {self.config.username}\033[0m\n\nTable Schema:\n"
width = 25
fields = 3
_repr += "-" * (width * fields + 1) + "\n"
columns = ["name", "type", "key"]
_repr += f"|\033[94m{columns[0]:24s}\033[0m|\033[96m{columns[1]:24s}"
_repr += f"\033[0m|\033[96m{columns[2]:24s}\033[0m|\n"
_repr += "-" * (width * fields + 1) + "\n"
q_str = f"DESC {self.config.database}.{self.config.table}"
debug_output(q_str)
rs = get_named_result(self.connection, q_str)
for r in rs:
_repr += f"|\033[94m{r['Field']:24s}\033[0m|\033[96m{r['Type']:24s}"
_repr += f"\033[0m|\033[96m{r['Key']:24s}\033[0m|\n"
_repr += "-" * (width * fields + 1) + "\n"
return _repr
def _build_query_sql(
self, q_emb: List[float], topk: int, where_str: Optional[str] = None
) -> str:
q_emb_str = ",".join(map(str, q_emb))
if where_str:
where_str = f"WHERE {where_str}"
else:
where_str = ""
q_str = f"""
SELECT {self.config.column_map['document']},
{self.config.column_map['metadata']},
cosine_similarity_norm(array<float>[{q_emb_str}],
{self.config.column_map['embedding']}) as dist
FROM {self.config.database}.{self.config.table}
{where_str}
ORDER BY dist {self.dist_order}
LIMIT {topk}
"""
debug_output(q_str)
return q_str
[docs] def similarity_search(
self, query: str, k: int = 4, where_str: Optional[str] = None, **kwargs: Any
) -> List[Document]:
"""使用StarRocks进行相似性搜索
参数:
query (str): 查询字符串
k (int, optional): 要检索的前K个邻居。默认为4。
where_str (Optional[str], optional): where条件字符串。
默认为None。
注意: 请不要让最终用户填写此内容,并始终注意SQL注入问题。
在处理元数据时,请记住使用`{self.metadata_column}.attribute`而不是仅使用`attribute`。
其默认名称为`metadata`。
返回:
List[Document]: 文档列表
"""
return self.similarity_search_by_vector(
self.embedding_function.embed_query(query), k, where_str, **kwargs
)
[docs] def similarity_search_by_vector(
self,
embedding: List[float],
k: int = 4,
where_str: Optional[str] = None,
**kwargs: Any,
) -> List[Document]:
"""使用StarRocks通过向量执行相似性搜索
参数:
query (str): 查询字符串
k (int, optional): 要检索的前K个邻居。默认为4。
where_str (Optional[str], optional): where条件字符串。
默认为None。
注意: 请不要让最终用户填写此内容,并始终注意SQL注入问题。
处理元数据时,请记得使用`{self.metadata_column}.attribute`而不是仅使用`attribute`。
其默认名称为`metadata`。
返回:
List[Document]: (Document, 相似度)列表
"""
q_str = self._build_query_sql(embedding, k, where_str)
try:
return [
Document(
page_content=r[self.config.column_map["document"]],
metadata=json.loads(r[self.config.column_map["metadata"]]),
)
for r in get_named_result(self.connection, q_str)
]
except Exception as e:
logger.error(f"\033[91m\033[1m{type(e)}\033[0m \033[95m{str(e)}\033[0m")
return []
[docs] def similarity_search_with_relevance_scores(
self, query: str, k: int = 4, where_str: Optional[str] = None, **kwargs: Any
) -> List[Tuple[Document, float]]:
"""使用StarRocks执行相似性搜索
参数:
query (str): 查询字符串
k (int, optional): 要检索的前K个邻居。默认为4。
where_str (Optional[str], optional): where条件字符串。
默认为None。
注意: 请不要让最终用户填写此内容,并始终注意SQL注入问题。
处理元数据时,请记得使用`{self.metadata_column}.attribute`而不是仅使用`attribute`。
其默认名称为`metadata`。
返回:
List[Document]: 文档列表
"""
q_str = self._build_query_sql(
self.embedding_function.embed_query(query), k, where_str
)
try:
return [
(
Document(
page_content=r[self.config.column_map["document"]],
metadata=json.loads(r[self.config.column_map["metadata"]]),
),
r["dist"],
)
for r in get_named_result(self.connection, q_str)
]
except Exception as e:
logger.error(f"\033[91m\033[1m{type(e)}\033[0m \033[95m{str(e)}\033[0m")
return []
[docs] def drop(self) -> None:
"""
辅助函数:丢弃数据
"""
get_named_result(
self.connection,
f"DROP TABLE IF EXISTS {self.config.database}.{self.config.table}",
)
@property
def metadata_column(self) -> str:
return self.config.column_map["metadata"]