Untitled
unknown
plain_text
a year ago
2.9 kB
7
Indexable
import boto3 from pydantic import BaseModel from typing import List import json class UserPrompt(BaseModel): query: str class BusinessUnitSearchResultItem(BaseModel): lookup_value: str linked_records: List[str] class BusinessUnitSearchResults(BaseModel): hits: List[BusinessUnitSearchResultItem] found: int class BusinessUnitSearchParameters(BaseModel): query: str fields: List[str] = ["lookup_value"] sort: str = '_score desc' size: int = 10 class BusinessUnitCloudSearchClient: def __init__(self, domain_endpoint: str): self.client = boto3.client('cloudsearchdomain', endpoint_url=domain_endpoint) def search(self, params: BusinessUnitSearchParameters) -> BusinessUnitSearchResults: response = self.client.search( query=params.query, queryParser='simple', sort=params.sort, size=params.size, queryOptions = json.dumps({"fields": params.fields}) ) search_results = [] for hit in response['hits']['hit']: search_results.append(BusinessUnitSearchResultItem( lookup_value=hit['fields'].get('lookup_value', [''])[0], linked_records=hit['fields'].get('linked_records', []) )) return BusinessUnitSearchResults( hits=search_results, found=response['hits']['found'] ) def handle_user_prompt_for_BusinessUnit(user_prompt: UserPrompt, client: BusinessUnitCloudSearchClient) -> BusinessUnitSearchResults: params = BusinessUnitSearchParameters(query=user_prompt.query) results = client.search(params) return results def format_business_unit_search_results(business_unit_search_results: BusinessUnitSearchResults, keyword: str) -> List[dict]: formatted_results = [] for item in business_unit_search_results.hits: matches = item.linked_records # Format matches list formatted_result = { "type": "business_unit", "index_name": "fincopilot-dim-business-unit", "total_count": business_unit_search_results.found, "extracted_count": len(matches), "matched_on": keyword, "lookup_value": item.lookup_value, "matches": matches } formatted_results.append(formatted_result) return formatted_results if __name__ == "__main__": with open('config.json', 'r') as f: config = json.load(f) client = BusinessUnitCloudSearchClient(config['BUSINESS_UNIT_DOMAIN_ENDPOINT']) query="Australia" user_prompt = UserPrompt(query="marketplace") results = handle_user_prompt_for_BusinessUnit(user_prompt=user_prompt, client=client) formatted_results = format_business_unit_search_results(results, query) print(formatted_results)
Editor is loading...
Leave a Comment