anthony/ollama
1
0
Fork 0
Code Activity

Много мелких доработок

Browse Source 
- переименован input_md => data
- добавление инфы о дате, версии и авторе изменений conf-страницы в индекс
- вывод этой инфы в источниках
- вывод статистики последнего ответа
- указание имени коллекции для qdrant
- мелочи по текстовкам
This commit is contained in:
Anthony Axenov
2025-08-29 08:54:43 +08:00
parent 3f2491db27
commit f3672e6ffd
10 changed files with 136 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
rag/rag.py
View File

@@ -6,7 +6,7 @@ from sentence_transformers import SentenceTransformer
class RagSystem:
def __init__(self,
md_folder: str = "input_md",
md_folder: str = "data",
ollama_url: str = "http://localhost:11434",
qdrant_host: str = "localhost",
qdrant_port: int = 6333,
@@ -27,13 +27,14 @@ class RagSystem:
body = {"model": self.chat_model}
requests.post(url, json=body, timeout=600)
def search_qdrant(self, query: str, top_k: int = 6):
def search_qdrant(self, query: str, top_k: int = 6, qdrant_collection="rag"):
query_vec = self.emb_model.encode(query, show_progress_bar=False).tolist()
url = f"http://{self.qdrant_host}:{self.qdrant_port}/collections/rag_collection/points/search"
url = f"http://{self.qdrant_host}:{self.qdrant_port}/collections/{qdrant_collection}/points/search"
payload = {
"vector": query_vec,
"top": top_k,
"with_payload": True
"with_payload": True,
# "score_threshold": 0.6
}
resp = requests.post(url, json=payload)
if resp.status_code != 200:
@@ -114,7 +115,7 @@ Respond."""
# "top_p": 0.1,
# },
}
self.response = requests.post(url, json=body, timeout=600)
self.response = requests.post(url, json=body, timeout=900)
if self.response.status_code != 200:
return f"Ошибка генерации ответа: {self.response.status_code} {self.response.text}"
return self.response.json().get("response", "").strip()
@@ -127,7 +128,7 @@ Respond."""
"messages": self.conversation_history,
"stream": True
}
resp = requests.post(url, json=body, stream=True, timeout=600)
resp = requests.post(url, json=body, stream=True, timeout=900)
if resp.status_code != 200:
raise RuntimeError(f"Ошибка генерации ответа: {resp.status_code} {resp.text}")
full_answer = ""
@@ -148,18 +149,28 @@ Respond."""
print(f"Error processing chunk: {e}")
def get_prompt_eval_count(self):
if not self.response:
return 0
return self.response.json().get("prompt_eval_count", 0)
def get_prompt_eval_duration(self):
if not self.response:
return 0
return self.response.json().get("prompt_eval_duration", 0) / (10 ** 9)
def get_eval_count(self):
if not self.response:
return 0
return self.response.json().get("eval_count", 0)
def get_eval_duration(self):
if not self.response:
return 0
return self.response.json().get("eval_duration", 0) / (10 ** 9)
def get_total_duration(self):
if not self.response:
return 0
return self.response.json().get("total_duration", 0) / (10 ** 9)
def get_tps(self):
@@ -172,19 +183,31 @@ Respond."""
def print_sources(context_docs: list):
print("\n\nИсточники:")
for idx, doc in enumerate(context_docs, start=1):
filename = doc['payload'].get("filename", None)
title = doc['payload'].get("filename", None)
url = doc['payload'].get("url", None)
title = filename
date = doc['payload'].get("date", None)
version = doc['payload'].get("version", None)
author = doc['payload'].get("author", None)
if url is None:
url = "(нет веб-ссылки)"
print(f"{idx}. {title}\n {url}\n")
if date is None:
date = "(неизвестно)"
if version is None:
version = "0"
if author is None:
author = "(неизвестен)"
print(f"{idx}. {title}")
print(f" {url} (v{version} {author})")
print(f" актуальность на {date}")
def print_v(text: str, is_verbose: bool):
if is_verbose:
print(text)
def print_stats(rag: RagSystem):
print("Статистика:")
print("\n\nСтатистика:")
print(f"* Time: {rag.get_total_duration()}s")
print(f"* TPS: {rag.get_tps()}")
print(f"* PEC: {rag.get_prompt_eval_count()}")
@@ -202,6 +225,7 @@ def main():
parser.add_argument("--show-prompt", default=False, action=argparse.BooleanOptionalAction, help="Показать полный промпт перед обработкой запроса")
parser.add_argument("--qdrant-host", default="localhost", help="Qdrant host")
parser.add_argument("--qdrant-port", type=int, default=6333, help="Qdrant port")
parser.add_argument("--qdrant-collection", type=str, default="rag", help="Название коллекции для поиска документов")
parser.add_argument("--ollama-url", default="http://localhost:11434", help="Ollama API URL")
parser.add_argument("--emb-model", default="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", help="Модель эмбеддинга")
parser.add_argument("--chat-model", default="phi4-mini:3.8b", help="Модель генерации Ollama")
@@ -220,6 +244,7 @@ def main():
print_v(f"Модель эмбеддинга: {args.emb_model}", args.verbose)
print_v(f"Модель чата: {args.chat_model}", args.verbose)
print_v(f"Документов для поиска: {args.topk}", args.verbose)
print_v(f"Коллекция для поиска: {args.qdrant_collection}", args.verbose)
if os.path.exists('sys_prompt.txt'):
print_v("Будет использоваться sys_prompt.txt!", args.verbose)
@@ -252,8 +277,9 @@ def main():
if query.lower() == "help":
print("<<< Команды итерактивного режима:")
print("save -- сохранить диалог в файл")
print("exit -- выход\n")
print("save -- сохранить диалог в файл")
print("stats -- вывести статистику последнего ответа")
print("exit -- выход\n")
query = None
continue
@@ -290,9 +316,9 @@ def main():
break
print_v("\nПоиск релевантных документов...", args.verbose)
context_docs = rag.search_qdrant(query, top_k=args.topk)
context_docs = rag.search_qdrant(query, top_k=args.topk, qdrant_collection=args.qdrant_collection)
if not context_docs:
print_v("Релевантные документы не найдены.", args.verbose)
print("<<< Релевантные документы не найдены")
if args.interactive:
query = None
continue