import boto3
import json
from opensearchpy import OpenSearch

bedrock = boto3.client("bedrock-runtime")

def get_embedding(text):
    body = json.dumps({"inputText": text})
    response = bedrock.invoke_model(
        modelId="amazon.titan-embed-text-v1",
        body=body
    )
    result = json.loads(response["body"].read())
    return result["embedding"]

client = OpenSearch(
    hosts=[{"host": "YOUR-ENDPOINT", "port": 443}]
)

query = "What is Python list?"
query_vector = get_embedding(query)

search_body = {
    "size": 1,
    "query": {
        "knn": {
            "embedding": {
                "vector": query_vector,
                "k": 1
            }
        }
    }
}

response = client.search(
    index="rag-index",
    body=search_body
)

context = response["hits"]["hits"][0]["_source"]["text"]

prompt = f"Answer based on context: {context}\nQuestion: {query}"

llm_response = bedrock.invoke_model(
    modelId="amazon.titan-text-lite-v1",
    body=json.dumps({"inputText": prompt})
)

answer = json.loads(llm_response["body"].read())

print(answer)