fix

test.py

@@ -1,28 +1,13 @@
-import pandas as pd
-from translate import Translator
+import asyncio
 
+import numpy as np
 
-def convert_and_translate_headers(input_file: str, output_file: str, sheet_name: str = None):
-    try:
-        # Читаем файл Excel
-        data = pd.read_excel(input_file, sheet_name=sheet_name)
+from trauma.api.message.ai.openai_request import convert_value_to_embeddings
+from trauma.core.config import settings
 
-        # Инициализируем переводчик
-        translator = Translator(from_lang='nl', to_lang='en')
 
-        # Переводим названия колонок
-        translated_columns = {col: translator.translate(col) for col in data.columns}
-        data.rename(columns=translated_columns, inplace=True)
+async def main():
+    entities = await settings.DB_CLIENT
 
-        # Сохраняем преобразованные данные в CSV
-        data.to_csv(output_file, index=False)
-        print(f"Файл успешно конвертирован и сохранен: {output_file}")
-    except Exception as e:
-        print(f"Произошла ошибка: {e}")
-
-
-input_xlsx = "test.xlsx"  # Путь к входному .xlsx файлу
-output_csv = "translated_output.csv"  # Путь к выходному .csv файлу
-sheet = "Sheet1"  # Укажите имя листа, если нужно
-
-convert_and_translate_headers(input_xlsx, output_csv, sheet)
+if __name__ == '__main__':
+    asyncio.run(main())

trauma/api/data/data/prepare_data.py

@@ -130,5 +130,27 @@ async def generate_descriptions():
         )
 
 
+def split_array(array, max_size=10):
+    if max_size <= 0:
+        raise ValueError("max_size must be greater than 0")
+    return [array[i:i + max_size] for i in range(0, len(array), max_size)]
+
+
+async def generate_embeddings():
+    entities = await settings.DB_CLIENT.entities.find({}, {"embedding": 0}).to_list()
+    entities = [EntityModel.from_mongo(entity) for entity in entities]
+    for entity in entities:
+        entity_str = entity.name
+        entity_emb = await settings.OPENAI_CLIENT.embeddings.create(
+            model='text-embedding-3-large',
+            dimensions=384,
+            input=entity_str,
+        )
+        await settings.DB_CLIENT.entities.update_one(
+            {"id": entity.id},
+            {"$set": {"embedding": entity_emb.data[0].embedding}},
+        )
+    print('hi')
+
 if __name__ == '__main__':
-    asyncio.run(generate_descriptions())
+    asyncio.run(generate_embeddings())

trauma/api/data/db_requests.py

@@ -5,7 +5,7 @@ from trauma.core.config import settings
 
 
 async def get_facility_by_id(facility_id: str) -> EntityModel:
-    facility = await settings.DB_CLIENT.entities.find_one({"id": facility_id})
+    facility = await settings.DB_CLIENT.entities.find_one({"id": facility_id}, {"embedding": 0})
     if not facility:
         raise HTTPException(status_code=404, detail="Country with specified id doesn't exists.")
     return EntityModel.from_mongo(facility)
@@ -13,5 +13,5 @@ async def get_facility_by_id(facility_id: str) -> EntityModel:
 
 async def get_all_model_obj() -> list[EntityModel]:
     sort_v = -1
-    objects = await settings.DB_CLIENT.entities.find().sort("_id", sort_v).to_list(length=None)
+    objects = await settings.DB_CLIENT.entities.find({}, {"embedding": 0}).sort("_id", sort_v).to_list(length=None)
     return objects

trauma/api/data/model.py

@@ -13,8 +13,8 @@ class EntityModel(MongoBaseModel):
 
 
 class EntityModelExtended(EntityModel):
-    highlightedAgeGroup: AgeGroup
-    highlightedTreatmentArea: str | None
-    highlightedTreatmentMethod: str | None
+    highlightedAgeGroup: AgeGroup | None = None
+    highlightedTreatmentArea: str | None = None
+    highlightedTreatmentMethod: str | None = None
     topMatch: bool = False
     score: float | None = None

trauma/api/message/ai/engine.py

@@ -10,7 +10,8 @@ from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
                                                   generate_search_request,
                                                   generate_final_response, convert_value_to_embeddings,
                                                   choose_closest_treatment_method, choose_closest_treatment_area,
-                                                  check_is_valid_request, generate_invalid_response, set_entity_score)
+                                                  check_is_valid_request, generate_invalid_response, set_entity_score,
+                                                  retrieve_semantic_answer, generate_searched_entity_response)
 from trauma.api.message.db_requests import (save_assistant_user_message,
                                             filter_entities_by_age_location,
                                             update_entity_data_obj, get_entity_by_index)
@@ -25,10 +26,16 @@ from trauma.core.config import settings
 async def search_entities(
         user_message: str, messages: list[dict], chat: ChatModel
 ) -> CreateMessageResponse:
-    entity_data, is_valid = await asyncio.gather(
+    related_entity, entity_data, is_valid = await asyncio.gather(
+        retrieve_semantic_answer(user_message),
         update_entity_data_with_ai(chat.entityData, user_message, messages[-1]['content']),
         check_is_valid_request(user_message, "\n".join([f"- [{i['role']}]: {i['content']}." for i in messages]))
     )
+    if related_entity:
+        response = await generate_searched_entity_response(user_message, related_entity[0])
+        asyncio.create_task(save_assistant_user_message(user_message, response, chat.id))
+        return CreateMessageResponse(text=response, entities=related_entity)
+
     final_entities = None
     if not is_valid:
         response = await generate_invalid_response(user_message, messages)

trauma/api/message/ai/openai_request.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 
 from trauma.api.chat.dto import EntityData
@@ -73,11 +74,11 @@ async def generate_final_response(final_entities: str, user_message: str, messag
     return messages
 
 
-async def convert_value_to_embeddings(value: str) -> list[float]:
+async def convert_value_to_embeddings(value: str, dimensions: int = 1536) -> list[float]:
     embeddings = await settings.OPENAI_CLIENT.embeddings.create(
         input=value,
         model='text-embedding-3-large',
-        dimensions=1536,
+        dimensions=dimensions,
     )
     return embeddings.data[0].embedding
 
@@ -150,3 +151,40 @@ async def set_entity_score(entity: EntityModelExtended, search_request: str):
         }
     ]
     return messages
+
+
+async def retrieve_semantic_answer(user_query: str) -> list[EntityModelExtended] | None:
+    embedding = await settings.OPENAI_CLIENT.embeddings.create(input=user_query,
+                                                               model='text-embedding-3-large',
+                                                               dimensions=384)
+    response = await settings.DB_CLIENT.entities.aggregate([
+        {"$vectorSearch": {
+            "index": f"entityVectors",
+            "path": "embedding",
+            "queryVector": embedding.data[0].embedding,
+            "numCandidates": 20,
+            "limit": 1
+        }},
+        {"$project": {
+            "embedding": 0,
+            "score": {"$meta": "vectorSearchScore"}
+        }}
+    ]).to_list(length=1)
+    return [EntityModelExtended(**response[0])] if response[0]['score'] > 0.83 else None
+
+
+@openai_wrapper()
+async def generate_searched_entity_response(user_query: str, facility: EntityModelExtended):
+    messages = [
+        {
+            "role": "system",
+            "content": TraumaPrompts.generate_searched_entity
+            .replace("{user_query}", user_query)
+            .replace("{entity}", facility.model_dump_json(indent=2))
+        }
+    ]
+    return messages
+
+
+if __name__ == '__main__':
+    asyncio.run(retrieve_semantic_answer('I want to know more about Praktijk Hermens'))

trauma/api/message/ai/prompts.py

@@ -167,6 +167,25 @@ De gebruiker zoekt naar een geschikte kliniek voor een patiënt en deelt hierbij
 - Gebruik een vriendelijke en geruststellende toon, bijvoorbeeld: "Ik heb op basis van de ingevoerde gegevens geen kliniek kunnen vinden. Kunnen we samen kijken of we de informatie iets kunnen aanpassen om betere resultaten te krijgen?"
 - Geef praktische suggesties, zoals: "Misschien helpt het om meer details over de locatie of de behandelmethode te delen."
 - Stel open vragen om de gebruiker te begeleiden bij het verfijnen van de gegevens, zoals: "Zijn er andere belangrijke punten die we kunnen toevoegen?"""
+    generate_searched_entity = """## Taak
+
+Je moet de gevraagde faciliteit beschrijven, waarvan de informatie wordt gegeven in de sectie `## Data`. Analyseer de gebruikersvraag en de informatie over de faciliteit, en geef een beknopt en bondig antwoord.
+
+## Gegevens
+
+**Gebruikersvraag**:
+```
+{user_query}
+```
+
+**Faciliteit**:
+```
+{entity}
+```
+
+## Belangrijke opmerkingen
+
+- Je antwoord moet beknopt zijn."""
 
     generate_clinic_description = """## Taak
 

trauma/api/message/db_requests.py

@@ -80,7 +80,7 @@ async def filter_entities_by_age_location(entity_data: dict) -> list[int]:
             "$regex": f".*{entity_data['postalCode']}.*",
             "$options": "i"
         }
-    entities = await settings.DB_CLIENT.entities.find(query).to_list(length=None)
+    entities = await settings.DB_CLIENT.entities.find(query, {"embedding": 0}).to_list(length=None)
     return [entity['index'] for entity in entities]
 
 

trauma/api/message/utils.py

@@ -1,6 +1,5 @@
 import json
 
-from trauma.api.chat.dto import EntityData
 from trauma.api.data.dto import AgeGroup
 from trauma.api.data.model import EntityModel
 from trauma.api.message.model import MessageModel