| | import asyncio |
| | import json |
| | import re |
| | from typing import List, Dict |
| | import faiss |
| | import httpx |
| | import numpy as np |
| | import pandas as pd |
| | from sqlalchemy.ext.asyncio import AsyncSession |
| | from starlette.websockets import WebSocket |
| | from transformers import pipeline |
| |
|
| | from project.bot.models import MessagePair |
| | from project.config import settings |
| |
|
| |
|
| | class SearchBot: |
| | chat_history = [] |
| |
|
| | |
| | |
| |
|
| | def __init__(self, memory=None): |
| | if memory is None: |
| | memory = [] |
| | self.chat_history = memory |
| |
|
| | @staticmethod |
| | def _cls_pooling(model_output): |
| | return model_output.last_hidden_state[:, 0] |
| |
|
| | @staticmethod |
| | async def enrich_information_from_google(search_word: str) -> str: |
| | url = "https://places.googleapis.com/v1/places:searchText" |
| | headers = { |
| | "Content-Type": "application/json", |
| | "X-Goog-Api-Key": settings.GOOGLE_PLACES_API_KEY, |
| | "X-Goog-FieldMask": "places.shortFormattedAddress,places.websiteUri,places.internationalPhoneNumber," |
| | "places.googleMapsUri,places.photos" |
| | } |
| | data = { |
| | "textQuery": f"{search_word} in Javea", |
| | "languageCode": "nl", |
| | "maxResultCount": 1, |
| |
|
| | } |
| | async with httpx.AsyncClient() as client: |
| | response = await client.post(url, headers=headers, content=json.dumps(data)) |
| | place_response = response.json() |
| | place_response = place_response['places'][0] |
| | photo_name = place_response.get('photos') |
| | photo_uri = None |
| | if photo_name: |
| | async with httpx.AsyncClient() as client: |
| | response = await client.get( |
| | f'https://places.googleapis.com/v1/{photo_name[0]["name"]}/media?maxWidthPx=350&key={settings.GOOGLE_PLACES_API_KEY}') |
| | photo_response = response.json() |
| | photo_uri = photo_response.get('photoUri') |
| | google_maps_uri = place_response.get('googleMapsUri') |
| | phone_number = place_response.get('internationalPhoneNumber') |
| | formatted_address = place_response.get('shortFormattedAddress') |
| | website_uri = place_response.get('websiteUri') |
| | if not google_maps_uri: |
| | return search_word |
| | enriched_word = f'<a class="extraDataLink" href="{google_maps_uri}" target="_blank">{search_word}</a><div class="tooltip-elem">' |
| | if photo_uri: |
| | enriched_word += f'<img src="{photo_uri}" alt="Image" class="tooltip-img">' |
| | if formatted_address: |
| | enriched_word += f'<p><a href="{google_maps_uri}" target="_blank">{formatted_address}</a></p>' |
| | if website_uri: |
| | enriched_word += f'<p><a href="{website_uri}">Google Maps URI</a></p>' |
| | if phone_number: |
| | phone_str = re.sub(r' ', '', phone_number) |
| | enriched_word += f'<p><a href="tel:{phone_str}">Phone number</a></p>' |
| | enriched_word += f"</div>" |
| | return enriched_word |
| |
|
| | async def analyze_full_response(self) -> str: |
| | assistant_message = self.chat_history.pop()['content'] |
| | nlp = pipeline("ner", model=settings.NLP_MODEL, tokenizer=settings.NLP_TOKENIZER, aggregation_strategy="simple") |
| | ner_result = nlp(assistant_message) |
| | analyzed_assistant_message = assistant_message |
| | for entity in ner_result: |
| | if entity['entity_group'] in ("LOC", "ORG", "MISC") and entity['word'] != "Javea": |
| | enriched_information = await self.enrich_information_from_google(entity['word']) |
| | analyzed_assistant_message = analyzed_assistant_message.replace(entity['word'], enriched_information, 1) |
| | return "ENRICHED:" + analyzed_assistant_message |
| |
|
| | async def _convert_to_embeddings(self, text_list): |
| | encoded_input = settings.INFO_TOKENIZER( |
| | text_list, padding=True, truncation=True, return_tensors="pt" |
| | ) |
| | encoded_input = {k: v.to(settings.device) for k, v in encoded_input.items()} |
| | model_output = settings.INFO_MODEL(**encoded_input) |
| | return self._cls_pooling(model_output).cpu().detach().numpy().astype('float32') |
| |
|
| | @staticmethod |
| | async def _get_context_data(user_query: list[float]) -> list[dict]: |
| | radius = 5 |
| | _, distances, indices = settings.FAISS_INDEX.range_search(user_query, radius) |
| | indices_distances_df = pd.DataFrame({'index': indices, 'distance': distances}) |
| | filtered_data_df = settings.products_dataset.iloc[indices].copy() |
| | filtered_data_df.loc[:, 'distance'] = indices_distances_df['distance'].values |
| | sorted_data_df: pd.DataFrame = filtered_data_df.sort_values(by='distance').reset_index(drop=True) |
| | sorted_data_df = sorted_data_df.drop('distance', axis=1) |
| | data = sorted_data_df.head(3).to_dict(orient='records') |
| | cleaned_data = [] |
| | for chunk in data: |
| | if "Comments:" in chunk['chunks']: |
| | cleaned_data.append(chunk) |
| | return cleaned_data |
| |
|
| | @staticmethod |
| | async def create_context_str(context: List[Dict]) -> str: |
| | context_str = '' |
| | for i, chunk in enumerate(context): |
| | context_str += f'{i + 1}) {chunk["chunks"]}' |
| | return context_str |
| |
|
| | async def _rag(self, context: List[Dict], query: str, session: AsyncSession, country: str): |
| | if context: |
| | context_str = await self.create_context_str(context) |
| | assistant_message = {"role": 'assistant', "content": context_str} |
| | self.chat_history.append(assistant_message) |
| | content = settings.PROMPT |
| | else: |
| | content = settings.EMPTY_PROMPT |
| | user_message = {"role": 'user', "content": query} |
| | self.chat_history.append(user_message) |
| | messages = [ |
| | { |
| | 'role': 'system', |
| | 'content': content |
| | }, |
| | ] |
| | messages = messages + self.chat_history |
| |
|
| | stream = await settings.OPENAI_CLIENT.chat.completions.create( |
| | messages=messages, |
| | temperature=0.1, |
| | n=1, |
| | model="gpt-3.5-turbo", |
| | stream=True |
| | ) |
| | response = '' |
| | async for chunk in stream: |
| | if chunk.choices[0].delta.content is not None: |
| | chunk_content = chunk.choices[0].delta.content |
| | response += chunk_content |
| | yield response |
| | await asyncio.sleep(0.02) |
| | assistant_message = {"role": 'assistant', "content": response} |
| | self.chat_history.append(assistant_message) |
| | try: |
| | session.add(MessagePair(user_message=query, bot_response=response, country=country)) |
| | except Exception as e: |
| | print(e) |
| |
|
| | async def ask_and_send(self, data: Dict, websocket: WebSocket, session: AsyncSession): |
| | query = data['query'] |
| | country = data['country'] |
| | transformed_query = await self._convert_to_embeddings(query) |
| | context = await self._get_context_data(transformed_query) |
| | try: |
| | async for chunk in self._rag(context, query, session, country): |
| | await websocket.send_text(chunk) |
| | analyzing = await self.analyze_full_response() |
| | await websocket.send_text(analyzing) |
| | except Exception: |
| | await self.emergency_db_saving(session) |
| |
|
| | @staticmethod |
| | async def emergency_db_saving(session: AsyncSession): |
| | await session.commit() |
| | await session.close() |
| |
|
