Untitled

mail@pastecode.io avatar
unknown
plain_text
a month ago
3.7 kB
5
Indexable
Never
from pydantic import BaseModel

# Define your desired output structure
class UserInfo(BaseModel):
    account: str
    department: str
    period: str
    subsidiary: str
    vendor: str

def get_user_info(client, message: str) -> UserInfo:
    # Extract structured data from natural language
    user_info = client.chat.completions.create(
        model="gpt-3.5-turbo",
        response_model=UserInfo,
        messages=[{"role": "user", "content": message}],
    )
    return user_info


import boto3
from pydantic import BaseModel
from typing import List

class UserPrompt(BaseModel):
    query: str

class SearchResultItem(BaseModel):
    id: str
    title: str
    content: str
    score: str
    subsidiary_full_name: str
    child_subsidiary_details: str
    subsidiary_id: str
    subsidiary_key: str

class SearchResults(BaseModel):
    hits: List[SearchResultItem]
    found: int

class CloudSearchClient:
    def __init__(self, domain_endpoint: str):
        self.client = boto3.client('cloudsearchdomain', endpoint_url=domain_endpoint)

    def search(self, query: str) -> SearchResults:
        response = self.client.search(
            query=query,
            queryParser='simple',
            size=10
        )
        search_results = []
        for hit in response['hits']['hit']:
            search_results.append(SearchResultItem(
                id=hit['id'],
                title=hit['fields'].get('title', [''])[0],
                content=hit['fields'].get('content', [''])[0],
                score=hit['fields'].get('_score', [''])[0],
                subsidiary_full_name=hit['fields'].get('subsidiary_full_name', [''])[0],
                child_subsidiary_details=hit['fields'].get('child_subsidiary_details', [''])[0],
                subsidiary_id=hit['fields'].get('subsidiary_id', [''])[0],
                subsidiary_key=hit['fields'].get('subsidiary_key', [''])[0]
            ))

        return SearchResults(
            hits=search_results,
            found=response['hits']['found']
        )

# Configuration file to store API keys and CloudSearch endpoint
class Config:
    OPENAI_API_KEY = 'your_openai_api_key_here'
    CLOUDSEARCH_DOMAIN_ENDPOINT = 'your_cloudsearch_domain_endpoint_here'

def merge_child_subsidiary_details(search_results) -> dict:
    merged_details = {}
    for item in search_results.hits:
        key = item.subsidiary_key
        details = item.child_subsidiary_details
        if key not in merged_details:
            merged_details[key] = []
        merged_details[key].append(details)
    return merged_details

import instructor
from openai import OpenAI
from config import Config
from user_prompt_parser.py import get_user_info
from cloudsearch_client import CloudSearchClient, UserPrompt
from merging_util import merge_child_subsidiary_details

def main():
    # Initialize OpenAI client
    openai_client = instructor.from_openai(OpenAI(api_key=Config.OPENAI_API_KEY))

    # User message
    message = "Give me all bills & bill credits for the pre-paid rent account and subsidiary DDE-US for Q2?"

    # Parse user prompt
    user_info = get_user_info(openai_client, message)

    # CloudSearch integration
    user_prompt = UserPrompt(query=user_info.subsidiary)
    cloudsearch_client = CloudSearchClient(Config.CLOUDSEARCH_DOMAIN_ENDPOINT)
    search_results = cloudsearch_client.search(user_prompt.query)

    # Merge child subsidiary details
    merged_details = merge_child_subsidiary_details(search_results)

    # Print results (or handle them as needed)
    print(search_results.json(indent=2))
    print(json.dumps(merged_details, indent=2))

if __name__ == "__main__":
    main()
Leave a Comment