Upload 3 files

db.py

@@ -0,0 +1,63 @@
+import os
+import requests
+from chromadb import Client, Settings, PersistentClient
+from chromadb.utils.embedding_functions.sentence_transformer_embedding_function import SentenceTransformerEmbeddingFunction
+
+import logging
+
+logger = logging.getLogger("db")
+logging.basicConfig(
+    format="%(asctime)s %(levelname)-8s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+JINA_KEY = os.getenv('JINA_API_KEY')
+
+jina_headers = {
+"Authorization": f"Bearer {JINA_KEY}"
+}
+
+def get_data_url(url):
+    logger.info(f"Scraping {url}")
+    jina_response = requests.get(f"https://r.jina.ai/{url}", headers=jina_headers, verify=False)
+
+    return jina_response.text
+
+class HacatonDB:
+    def __init__(self):
+        self.client = PersistentClient(settings=Settings(anonymized_telemetry=False))
+        self.embed = SentenceTransformerEmbeddingFunction(
+            model_name="BAAI/bge-m3"
+        )
+        self.collection = self.client.create_collection('test_hakaton', embedding_function=self.embed, metadata={"hnsw:space": "cosine"}, get_or_create=True)
+
+    def add(self, urls):
+        logger.info(f"Add info to collection")
+        texts = []
+        meta = []
+        new_urls = []
+        for url in urls:
+            if len(self.collection.get(ids=[url])["ids"]) > 0:
+                logger.info(f"URL {url} already exist")
+                continue
+
+            new_urls.append(url)
+            texts.append(get_data_url(url))
+            meta.append({"file_name": f"file_{url.split('/')[-2]}"})
+            logger.info(f"URL {url} added")
+
+        if len(new_urls) > 0:
+            self.collection.add(documents=texts, ids=new_urls, metadatas=meta)
+            logger.info(f"Addition {len(new_urls)} sources completed")
+        else:
+            logger.info(f"No new sources")
+
+    def update(self, urls):
+        pass
+
+    def query(self, query, top_k):
+        return self.collection.query(query_texts=query, n_results=top_k)
+
+
+db = HacatonDB()

rag.py

@@ -0,0 +1,72 @@
+import os
+import openai
+from db import db
+import logging
+
+logger = logging.getLogger("rag")
+logging.basicConfig(
+    format="%(asctime)s %(levelname)-8s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+
+MISTRAL_KEY=os.getenv('MISTRAL_API_KEY')
+MISTRAL_URL="https://api.mistral.ai/v1"
+MISTRAL_MODEL="mistral-small-latest"
+
+
+client = openai.OpenAI(api_key=MISTRAL_KEY, base_url=MISTRAL_URL)
+
+message_template = """\
+Далее представлена информацию по опыту нашей компании
+---------------------
+{retrieved_chunks}
+---------------------
+Далее представлен запрос потенциального проекта
+---------------------
+{request_content}
+---------------------
+При проведении анализа опирайся только на представленную информацию"""
+
+# Функция для обработки запроса к LLM
+def process_query(req_file, system_prompt):
+    logger.info("Process query")
+
+    if req_file is not None:
+        with open(req_file.name, 'r', encoding='utf-8') as f:
+            req_file_content = f.read()
+    else:
+        logger.warning("File is not loaded!")
+        req_file_content = ""
+
+    logger.info("Retrive docs")
+
+    docs = db.query(req_file_content, top_k=1)
+
+    
+    logger.info(f"Retrived {len(docs['ids'][0])} docs")
+
+    doc_context = '\n\n'.join(docs['documents'][0])
+
+    # Создание контекста из файлов
+    user_message = message_template.format(retrieved_chunks=doc_context, request_content=req_file_content)
+
+    # Формирование сообщения для LLM
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_message}
+    ]
+
+    logger.info("LLM call")
+
+    response = client.chat.completions.create(
+            messages=messages,
+            model=MISTRAL_MODEL,
+        )
+    
+    logger.info("LLM call completed")
+
+    # Получение ответа от LLM
+    llm_response = response.choices[0].message.content
+    return llm_response

sources.txt

@@ -0,0 +1,3 @@
+https://www.reksoft.ru/blog/portfolio/iserver-ea-otp-bank/
+https://www.reksoft.ru/blog/portfolio/iserver-otkrytie-broker/
+https://www.reksoft.ru/blog/portfolio/func_test/