aphilippov/python-server-api
0
1from typing import List, Union, Callable2import os3import requests4from urllib.parse import urlparse5import glob6import chromadb7 8if chromadb.__version__ < "0.4.15":9 from chromadb.api import API10else:11 from chromadb.api import ClientAPI as API12from chromadb.api.types import QueryResult13import chromadb.utils.embedding_functions as ef14import logging15import pypdf16from autogen.token_count_utils import count_token17 18try:19 from unstructured.partition.auto import partition20 21 HAS_UNSTRUCTURED = True22except ImportError:23 HAS_UNSTRUCTURED = False24 25logger = logging.getLogger(__name__)26TEXT_FORMATS = [27 "txt",28 "json",29 "csv",30 "tsv",31 "md",32 "html",33 "htm",34 "rtf",35 "rst",36 "jsonl",37 "log",38 "xml",39 "yaml",40 "yml",41 "pdf",42]43UNSTRUCTURED_FORMATS = [44 "doc",45 "docx",46 "epub",47 "msg",48 "odt",49 "org",50 "pdf",51 "ppt",52 "pptx",53 "rtf",54 "rst",55 "xlsx",56] # These formats will be parsed by the 'unstructured' library, if installed.57if HAS_UNSTRUCTURED:58 TEXT_FORMATS += UNSTRUCTURED_FORMATS59 TEXT_FORMATS = list(set(TEXT_FORMATS))60VALID_CHUNK_MODES = frozenset({"one_line", "multi_lines"})61 62 63def split_text_to_chunks(64 text: str,65 max_tokens: int = 4000,66 chunk_mode: str = "multi_lines",67 must_break_at_empty_line: bool = True,68 overlap: int = 10,69):70 """Split a long text into chunks of max_tokens."""71 if chunk_mode not in VALID_CHUNK_MODES:72 raise AssertionError73 if chunk_mode == "one_line":74 must_break_at_empty_line = False75 chunks = []76 lines = text.split("\n")77 lines_tokens = [count_token(line) for line in lines]78 sum_tokens = sum(lines_tokens)79 while sum_tokens > max_tokens:80 if chunk_mode == "one_line":81 estimated_line_cut = 282 else:83 estimated_line_cut = int(max_tokens / sum_tokens * len(lines)) + 184 cnt = 085 prev = ""86 for cnt in reversed(range(estimated_line_cut)):87 if must_break_at_empty_line and lines[cnt].strip() != "":88 continue89 if sum(lines_tokens[:cnt]) <= max_tokens:90 prev = "\n".join(lines[:cnt])91 break92 if cnt == 0:93 logger.warning(94 f"max_tokens is too small to fit a single line of text. Breaking this line:\n\t{lines[0][:100]} ..."95 )96 if not must_break_at_empty_line:97 split_len = int(max_tokens / lines_tokens[0] * 0.9 * len(lines[0]))98 prev = lines[0][:split_len]99 lines[0] = lines[0][split_len:]100 lines_tokens[0] = count_token(lines[0])101 else:102 logger.warning("Failed to split docs with must_break_at_empty_line being True, set to False.")103 must_break_at_empty_line = False104 chunks.append(prev) if len(prev) > 10 else None # don't add chunks less than 10 characters105 lines = lines[cnt:]106 lines_tokens = lines_tokens[cnt:]107 sum_tokens = sum(lines_tokens)108 text_to_chunk = "\n".join(lines)109 chunks.append(text_to_chunk) if len(text_to_chunk) > 10 else None # don't add chunks less than 10 characters110 return chunks111 112 113def extract_text_from_pdf(file: str) -> str:114 """Extract text from PDF files"""115 text = ""116 with open(file, "rb") as f:117 reader = pypdf.PdfReader(f)118 if reader.is_encrypted: # Check if the PDF is encrypted119 try:120 reader.decrypt("")121 except pypdf.errors.FileNotDecryptedError as e:122 logger.warning(f"Could not decrypt PDF {file}, {e}")123 return text # Return empty text if PDF could not be decrypted124 125 for page_num in range(len(reader.pages)):126 page = reader.pages[page_num]127 text += page.extract_text()128 129 if not text.strip(): # Debugging line to check if text is empty130 logger.warning(f"Could not decrypt PDF {file}")131 132 return text133 134 135def split_files_to_chunks(136 files: list,137 max_tokens: int = 4000,138 chunk_mode: str = "multi_lines",139 must_break_at_empty_line: bool = True,140 custom_text_split_function: Callable = None,141):142 """Split a list of files into chunks of max_tokens."""143 144 chunks = []145 146 for file in files:147 _, file_extension = os.path.splitext(file)148 file_extension = file_extension.lower()149 150 if HAS_UNSTRUCTURED and file_extension[1:] in UNSTRUCTURED_FORMATS:151 text = partition(file)152 text = "\n".join([t.text for t in text]) if len(text) > 0 else ""153 elif file_extension == ".pdf":154 text = extract_text_from_pdf(file)155 else: # For non-PDF text-based files156 with open(file, "r", encoding="utf-8", errors="ignore") as f:157 text = f.read()158 159 if not text.strip(): # Debugging line to check if text is empty after reading160 logger.warning(f"No text available in file: {file}")161 continue # Skip to the next file if no text is available162 163 if custom_text_split_function is not None:164 chunks += custom_text_split_function(text)165 else:166 chunks += split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)167 168 return chunks169 170 171def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMATS, recursive: bool = True):172 """Return a list of all the files in a given directory, a url, a file path or a list of them."""173 if len(types) == 0:174 raise ValueError("types cannot be empty.")175 types = [t[1:].lower() if t.startswith(".") else t.lower() for t in set(types)]176 types += [t.upper() for t in types]177 178 files = []179 # If the path is a list of files or urls, process and return them180 if isinstance(dir_path, list):181 for item in dir_path:182 if os.path.isfile(item):183 files.append(item)184 elif is_url(item):185 files.append(get_file_from_url(item))186 elif os.path.exists(item):187 try:188 files.extend(get_files_from_dir(item, types, recursive))189 except ValueError:190 logger.warning(f"Directory {item} does not exist. Skipping.")191 else:192 logger.warning(f"File {item} does not exist. Skipping.")193 return files194 195 # If the path is a file, return it196 if os.path.isfile(dir_path):197 return [dir_path]198 199 # If the path is a url, download it and return the downloaded file200 if is_url(dir_path):201 return [get_file_from_url(dir_path)]202 203 if os.path.exists(dir_path):204 for type in types:205 if recursive:206 files += glob.glob(os.path.join(dir_path, f"**/*.{type}"), recursive=True)207 else:208 files += glob.glob(os.path.join(dir_path, f"*.{type}"), recursive=False)209 else:210 logger.error(f"Directory {dir_path} does not exist.")211 raise ValueError(f"Directory {dir_path} does not exist.")212 return files213 214 215def get_file_from_url(url: str, save_path: str = None):216 """Download a file from a URL."""217 if save_path is None:218 os.makedirs("/tmp/chromadb", exist_ok=True)219 save_path = os.path.join("/tmp/chromadb", os.path.basename(url))220 else:221 os.makedirs(os.path.dirname(save_path), exist_ok=True)222 with requests.get(url, stream=True) as r:223 r.raise_for_status()224 with open(save_path, "wb") as f:225 for chunk in r.iter_content(chunk_size=8192):226 f.write(chunk)227 return save_path228 229 230def is_url(string: str):231 """Return True if the string is a valid URL."""232 try:233 result = urlparse(string)234 return all([result.scheme, result.netloc])235 except ValueError:236 return False237 238 239def create_vector_db_from_dir(240 dir_path: Union[str, List[str]],241 max_tokens: int = 4000,242 client: API = None,243 db_path: str = "/tmp/chromadb.db",244 collection_name: str = "all-my-documents",245 get_or_create: bool = False,246 chunk_mode: str = "multi_lines",247 must_break_at_empty_line: bool = True,248 embedding_model: str = "all-MiniLM-L6-v2",249 embedding_function: Callable = None,250 custom_text_split_function: Callable = None,251 custom_text_types: List[str] = TEXT_FORMATS,252 recursive: bool = True,253 extra_docs: bool = False,254) -> API:255 """Create a vector db from all the files in a given directory, the directory can also be a single file or a url to256 a single file. We support chromadb compatible APIs to create the vector db, this function is not required if257 you prepared your own vector db.258 259 Args:260 dir_path (Union[str, List[str]]): the path to the directory, file, url or a list of them.261 max_tokens (Optional, int): the maximum number of tokens per chunk. Default is 4000.262 client (Optional, API): the chromadb client. Default is None.263 db_path (Optional, str): the path to the chromadb. Default is "/tmp/chromadb.db".264 collection_name (Optional, str): the name of the collection. Default is "all-my-documents".265 get_or_create (Optional, bool): Whether to get or create the collection. Default is False. If True, the collection266 will be returned if it already exists. Will raise ValueError if the collection already exists and get_or_create is False.267 chunk_mode (Optional, str): the chunk mode. Default is "multi_lines".268 must_break_at_empty_line (Optional, bool): Whether to break at empty line. Default is True.269 embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if270 embedding_function is not None.271 embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with272 the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding273 functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.274 custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings.275 Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`.276 custom_text_types (Optional, List[str]): a list of file types to be processed. Default is TEXT_FORMATS.277 recursive (Optional, bool): whether to search documents recursively in the dir_path. Default is True.278 extra_docs (Optional, bool): whether to add more documents in the collection. Default is False279 Returns:280 API: the chromadb client.281 """282 if client is None:283 client = chromadb.PersistentClient(path=db_path)284 try:285 embedding_function = (286 ef.SentenceTransformerEmbeddingFunction(embedding_model)287 if embedding_function is None288 else embedding_function289 )290 collection = client.create_collection(291 collection_name,292 get_or_create=get_or_create,293 embedding_function=embedding_function,294 # https://github.com/nmslib/hnswlib#supported-distances295 # https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184296 # https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md297 metadata={"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}, # ip, l2, cosine298 )299 300 length = 0301 if extra_docs:302 length = len(collection.get()["ids"])303 304 if custom_text_split_function is not None:305 chunks = split_files_to_chunks(306 get_files_from_dir(dir_path, custom_text_types, recursive),307 custom_text_split_function=custom_text_split_function,308 )309 else:310 chunks = split_files_to_chunks(311 get_files_from_dir(dir_path, custom_text_types, recursive),312 max_tokens,313 chunk_mode,314 must_break_at_empty_line,315 )316 logger.info(f"Found {len(chunks)} chunks.")317 # Upsert in batch of 40000 or less if the total number of chunks is less than 40000318 for i in range(0, len(chunks), min(40000, len(chunks))):319 end_idx = i + min(40000, len(chunks) - i)320 collection.upsert(321 documents=chunks[i:end_idx],322 ids=[f"doc_{j+length}" for j in range(i, end_idx)], # unique for each doc323 )324 except ValueError as e:325 logger.warning(f"{e}")326 return client327 328 329def query_vector_db(330 query_texts: List[str],331 n_results: int = 10,332 client: API = None,333 db_path: str = "/tmp/chromadb.db",334 collection_name: str = "all-my-documents",335 search_string: str = "",336 embedding_model: str = "all-MiniLM-L6-v2",337 embedding_function: Callable = None,338) -> QueryResult:339 """Query a vector db. We support chromadb compatible APIs, it's not required if you prepared your own vector db340 and query function.341 342 Args:343 query_texts (List[str]): the list of strings which will be used to query the vector db.344 n_results (Optional, int): the number of results to return. Default is 10.345 client (Optional, API): the chromadb compatible client. Default is None, a chromadb client will be used.346 db_path (Optional, str): the path to the vector db. Default is "/tmp/chromadb.db".347 collection_name (Optional, str): the name of the collection. Default is "all-my-documents".348 search_string (Optional, str): the search string. Only docs that contain an exact match of this string will be retrieved. Default is "".349 embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if350 embedding_function is not None.351 embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with352 the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding353 functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.354 355 Returns:356 QueryResult: the query result. The format is:357 class QueryResult(TypedDict):358 ids: List[IDs]359 embeddings: Optional[List[List[Embedding]]]360 documents: Optional[List[List[Document]]]361 metadatas: Optional[List[List[Metadata]]]362 distances: Optional[List[List[float]]]363 """364 if client is None:365 client = chromadb.PersistentClient(path=db_path)366 # the collection's embedding function is always the default one, but we want to use the one we used to create the367 # collection. So we compute the embeddings ourselves and pass it to the query function.368 collection = client.get_collection(collection_name)369 embedding_function = (370 ef.SentenceTransformerEmbeddingFunction(embedding_model) if embedding_function is None else embedding_function371 )372 query_embeddings = embedding_function(query_texts)373 # Query/search n most similar results. You can also .get by id374 results = collection.query(375 query_embeddings=query_embeddings,376 n_results=n_results,377 where_document={"$contains": search_string} if search_string else None, # optional filter378 )379 return results380