Spaces:
Running
Running
File size: 7,286 Bytes
e931b70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
import hashlib
from datetime import datetime
from typing import List, Optional
import pandas as pd
from clickhouse_connect import get_client
from langchain.schema.embeddings import Embeddings
from langchain.vectorstores.myscale import MyScaleWithoutJSON, MyScaleSettings
from streamlit.runtime.uploaded_file_manager import UploadedFile
from backend.chat_bot.tools import parse_files, extract_embedding
from backend.construct.build_retriever_tool import create_retriever_tool
from logger import logger
class ChatBotKnowledgeTable:
def __init__(self, host, port, username, password,
embedding: Embeddings, parser_api_key: str, db="chat",
kb_table="private_kb", tool_table="private_tool") -> None:
super().__init__()
personal_files_schema_ = f"""
CREATE TABLE IF NOT EXISTS {db}.{kb_table}(
entity_id String,
file_name String,
text String,
user_id String,
created_by DateTime,
vector Array(Float32),
CONSTRAINT cons_vec_len CHECK length(vector) = 768,
VECTOR INDEX vidx vector TYPE MSTG('metric_type=Cosine')
) ENGINE = ReplacingMergeTree ORDER BY entity_id
"""
# `tool_name` represent private knowledge database name.
private_knowledge_base_schema_ = f"""
CREATE TABLE IF NOT EXISTS {db}.{tool_table}(
tool_id String,
tool_name String,
file_names Array(String),
user_id String,
created_by DateTime,
tool_description String
) ENGINE = ReplacingMergeTree ORDER BY tool_id
"""
self.personal_files_table = kb_table
self.private_knowledge_base_table = tool_table
config = MyScaleSettings(
host=host,
port=port,
username=username,
password=password,
database=db,
table=kb_table,
)
self.client = get_client(
host=config.host,
port=config.port,
username=config.username,
password=config.password,
)
self.client.command("SET allow_experimental_object_type=1")
self.client.command(personal_files_schema_)
self.client.command(private_knowledge_base_schema_)
self.parser_api_key = parser_api_key
self.vector_store = MyScaleWithoutJSON(
embedding=embedding,
config=config,
must_have_cols=["file_name", "text", "created_by"],
)
# List all files with given `user_id`
def list_files(self, user_id: str):
query = f"""
SELECT DISTINCT file_name, COUNT(entity_id) AS num_paragraph,
arrayMax(arrayMap(x->length(x), groupArray(text))) AS max_chars
FROM {self.vector_store.config.database}.{self.personal_files_table}
WHERE user_id = '{user_id}' GROUP BY file_name
"""
return [r for r in self.vector_store.client.query(query).named_results()]
# Parse and embedding files
def add_by_file(self, user_id, files: List[UploadedFile]):
data = parse_files(self.parser_api_key, user_id, files)
data = extract_embedding(self.vector_store.embeddings, data)
self.vector_store.client.insert_df(
table=self.personal_files_table,
df=pd.DataFrame(data),
database=self.vector_store.config.database,
)
# Remove all files and private_knowledge_bases with given `user_id`
def clear(self, user_id: str):
self.vector_store.client.command(
f"DELETE FROM {self.vector_store.config.database}.{self.personal_files_table} "
f"WHERE user_id='{user_id}'"
)
query = f"""DELETE FROM {self.vector_store.config.database}.{self.private_knowledge_base_table}
WHERE user_id = '{user_id}'"""
self.vector_store.client.command(query)
def create_private_knowledge_base(
self, user_id: str, tool_name: str, tool_description: str, files: Optional[List[str]] = None
):
self.vector_store.client.insert_df(
self.private_knowledge_base_table,
pd.DataFrame(
[
{
"tool_id": hashlib.sha256(
(user_id + tool_name).encode("utf-8")
).hexdigest(),
"tool_name": tool_name, # tool_name represent user's private knowledge base.
"file_names": files,
"user_id": user_id,
"created_by": datetime.now(),
"tool_description": tool_description,
}
]
),
database=self.vector_store.config.database,
)
# Show all private knowledge bases with given `user_id`
def list_private_knowledge_bases(self, user_id: str, private_knowledge_base=None):
extended_where = f"AND tool_name = '{private_knowledge_base}'" if private_knowledge_base else ""
query = f"""
SELECT tool_name, tool_description, length(file_names)
FROM {self.vector_store.config.database}.{self.private_knowledge_base_table}
WHERE user_id = '{user_id}' {extended_where}
"""
return [r for r in self.vector_store.client.query(query).named_results()]
def remove_private_knowledge_bases(self, user_id: str, private_knowledge_bases: List[str]):
unique_list = list(set(private_knowledge_bases))
unique_list = ",".join([f"'{t}'" for t in unique_list])
query = f"""DELETE FROM {self.vector_store.config.database}.{self.private_knowledge_base_table}
WHERE user_id = '{user_id}' AND tool_name IN [{unique_list}]"""
self.vector_store.client.command(query)
def as_retrieval_tools(self, user_id, tool_name=None):
logger.info(f"")
private_knowledge_bases = self.list_private_knowledge_bases(user_id=user_id, private_knowledge_base=tool_name)
retrievers = {}
for private_kb in private_knowledge_bases:
file_names_sql = f"""
SELECT arrayJoin(file_names) FROM (
SELECT file_names
FROM chat.private_tool
WHERE user_id = '{user_id}' AND tool_name = '{private_kb["tool_name"]}'
)
"""
logger.info(f"user_id is {user_id}, file_names_sql is {file_names_sql}")
res = self.client.query(file_names_sql)
file_names = []
for line in res.result_rows:
file_names.append(line[0])
file_names = ', '.join(f"'{item}'" for item in file_names)
logger.info(f"user_id is {user_id}, file_names is {file_names}")
retrievers[private_kb["tool_name"]] = create_retriever_tool(
self.vector_store.as_retriever(
search_kwargs={"where_str": f"user_id='{user_id}' AND file_name IN ({file_names})"},
),
tool_name=private_kb["tool_name"],
description=private_kb["tool_description"],
)
return retrievers
|