app.py

from urllib import request
from pydantic import BaseModel
from fastapi import FastAPI, Request, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
from ibm_watsonx_ai import APIClient
from ibm_watsonx_ai import Credentials
from ibm_watsonx_ai.foundation_models import Model, ModelInference
from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
from ibm_watsonx_ai.foundation_models.utils.enums import ModelTypes, DecodingMethods
from dotenv import load_dotenv

import os
import http.client #token
import json #token
import requests #chroma
import time #measuring the time

load_dotenv()  # take environment variables from .env.

# load the environment variables
# get the IAM API Key from the environment variable    
iam_api_key = os.environ["WATSONX_IAM_APIKEY"]
#print(f"iam api key: {iam_api_key}")

watsonx_project_id = os.environ["WATSONX_PROJECT_ID"]
api_url = os.environ["WATSONX_API_URL"]
model_id = os.environ["WATSONX_MODEL_ID"]


number_of_responses_for_genai = int(os.environ["NUMBER_OF_RESPONSES_FOR_GENAI"])
retriver_discovery_or_chroma = os.environ["RETRIVER"]
json_chroma_for_wd_project = os.environ["CHROMA_FOR_WD_PROJECT"]

retriver_chroma_url = os.environ["CHROMA_URL"]

generate_params = {
    GenParams.MAX_NEW_TOKENS: 1000,
    GenParams.DECODING_METHOD: "greedy",
    GenParams.REPETITION_PENALTY: 1,
    GenParams.STOP_SEQUENCES: ["\\n\\n"]
}

watsonx_project_id=os.environ["WATSONX_PROJECT_ID"]

credentials=Credentials(
    api_key = iam_api_key, #WX_API_KEY,
    url = api_url #IBM_CLOUD_URL
    )
#print(credentials)

client=APIClient(credentials)
client_project = APIClient(credentials, project_id = watsonx_project_id)

MODEL_ID = model_id #"ibm/granite-8b-code-instruct" #"mistralai/mixtral-8x7b-instruct-v01"

# Initialize the app and add CORS middleware
app = FastAPI()
#app.add_middleware(
#    CORSMiddleware,
#    allow_origins=['*'],
#    allow_credentials=True,
#    allow_methods=['*'],
#    allow_headers=["*"]
#)


# code that connects to chromaDB API service per project;
def retrive_chroma(query, project_id):


    #select the container to ping
    url = retriver_chroma_url

    # pick the url by project_id use-case

    print("in retriver")
    # connect to the container
    #print(query)
    # get results
    try:
        payload = json.dumps({
            "query": query,
            "no_results": number_of_responses_for_genai
        })
        headers = {
            "Content-Type":"application/json"
        }
        #print(payload)
        #print(headers)
        response = requests.request("POST", url, headers=headers, data=payload, verify=False)

        #print(response.text)

        result = response.text
    except Exception as e:
        return {'error': str(e)}

    return result 

# code that adds an original query to the received top 3 answers and sends the data to watsonx.ai with the custom prompt
def augment_chroma(query, response_retriver):
    #check if there are any answers, if no return empty
    #{"results":{"data":null,"distances":[[0.29149818420410156,0.6023381948471069]],"documents":[["Cleanses are a pack of specific juices that are particularly picked to help the customer with a specific health problem or improve their specific health category.","list the juices and lemonades in the order provided above. - Recommend drinking water or tea between juices or lemonades. - For the Gut Whisperer cleanse, start the first juice at 8 AM and the last lemonade at 8 PM. - The Gut Whisperer cleanse is a specialty cleanse; find it by searching \"gut whisperer well juicery\" and visiting the Well Juicery website. - For other questions, refer to the guide on the Well Juicery website for detailed instructions on consuming the cleanses."]],"embeddings":null,"ids":[["id19","id28"]],"included":["metadatas","documents","distances"],"metadatas":[[{"question":"**Cleanses:**"},{"question":"If asked about the order of taking the cleanses,"}]],"uris":null}}
    print("response from retriver")
    #print (response_retriver)
    if response_retriver is None:
        return
    
    # TODO select top three or n
    # Add the original query to the received top 3 answers

    # Extract the answer and source from the query result
    answer = []
    source = []
    i = 0
    response_json=json.loads(response_retriver)
    for passage in response_json['results']['documents'][0]:
        #print(passage)
        if i < number_of_responses_for_genai :
            text1 = passage#[i]
            #print(f"text1 index {i}")
            #print(text1)
            question_json = response_json['results']['metadatas'][0][i]
            text2 = question_json["question"]
            found_answer= text2 + text1
            answer.append(found_answer)
            source.append(response_json['results']['ids'][0][i])#passage['document_id'])
        i=i+1
        print(f"counter: {i}")

    # Return the answer, and source in a JSON response
    
    # TODO reprioritize using some strategy
    # TODO add query to the response
    print(answer)
    return {
        'answer': answer,
        'source': source
    }


#suggested filter
# sample output
# {'model_id': 'ibm/granite-8b-code-instruct', 'model_version': '1.1.0', 'created_at': '2024-10-31T19:14:17.242Z', 'results': [{'generated_text': '\n', 'generated_token_count': 1, 'input_token_count': 0, 'stop_reason': 'not_finished'}]}

async def event_stream(watonsx_model: Model, llm_input, citation):
    start_gen_time = time.perf_counter()
    count = 0
    # async for chunk in watonsx_model.generate_text_stream(prompt=llm_input, raw_response=True):
    for chunk in watonsx_model.generate_text_stream(prompt=llm_input, raw_response=True):
        json_data = json.dumps(chunk)
        print(chunk)

        if count == 0:
            duration = int((time.perf_counter() - start_gen_time)*1000)
            print(f"Time taken for FIRST TOKEN: {duration}")
        count += 1

        if citation is not None and chunk["results"][0]["stop_reason"] == "eos_token":
            duration = int((time.perf_counter() - start_gen_time)*1000)
            print(f"Time taken for last token: {duration}")
            chunk["results"][0]["generated_text"] += citation
            json_data = json.dumps(chunk)

        yield f"data: {json_data}\n\n"  # SSE format



@app.get("/")
async def root():
    return {"message": "Hello Mengalo Streaming"}

@app.post('/query-streamed')
async def stream_response(request:Request):
    payload_data = await request.json()
    prompt = payload_data["prompt"].strip()
    my_model = payload_data["model"].strip()
        
    try:         
        query = payload_data["query"].strip()
        #model = ModelInference(
        model = Model(
            model_id = my_model, 
            params = generate_params, 
            credentials = credentials,
            project_id = watsonx_project_id #WX_PROJECT_ID
        )
        ### start here

        print("in query method")
        #data = request.get_json()
        #if data is None:
        #    return {"error": "Invalid request, JSON expected"}, 400
        #query = data.get("query")
        #project = query = payload_data["project_id"].strip()
        #if query is None or project is None:
        #    return {"error": "Missing query or project"}, 400
        
        ### get chroma
        # select chroma vs watson
        start_retriver_time = time.perf_counter()
        match retriver_discovery_or_chroma:
            #case "discovery": 
            #    response_retriver = retrive(query, project)
            #    response_augmenter = augment(query, response_retriver)
            case "chroma":
                print("in case chroma")
                response_retriver = retrive_chroma(query, watsonx_project_id)
                response_augmenter = augment_chroma(query, response_retriver)
            # use just the text of the document placed in the filesystem
            case _:
                response_retriver = retrive_chroma(query, watsonx_project_id)
                response_augmenter = augment_chroma(query, response_retriver)

        stop_retriver_time = time.perf_counter()
        retriver_elapsed_time = int((stop_retriver_time - start_retriver_time)*1000)
        print(f"retrived in: {retriver_elapsed_time} ms")
        #print(response_augmenter)

        url = api_url
        i = 0
        responses = ""
        for i in range(len(response_augmenter["answer"])):

            responses = responses + f"answer {i+1}: "+ "".join(response_augmenter["answer"][i]) + "source: "+ response_augmenter["source"][i] +"\n"
        print("my responses**********")
        print(responses)
        genai_prompt = f"""You are a conversational AI assistant tasked with answering customer inquiries clearly, accurately, and in a concise, friendly tone suitable for a phone conversation. 
        You have access to a retrieval system with responses from a pre-approved database to support your answers.
        Here is the customer question:

        Customer Inquiry:
        {query}

        When responding, ensure you:
        - Address the question directly and focus only on the customer's specific concern.
        - Include key details, such as product names, descriptions, benefits, ingredients, or prices, if relevant.
        - Use a helpful, warm tone that feels supportive and directly addresses the customer's needs.

        Specific guidance:
        - If asked about a product, provide ingredients, calorie content, and key benefits.
        - For questions on payments or shipping, offer clear instructions.
        - For questions about the company, briefly share values or product categories.
        - For complaints or issues (e.g., expired products), offer solutions like refunds, replacements, or direct escalation to a specialist when needed."""

        genai_prompt2 = f"""Refer to the following knowledge retrieval responses to help answer the customer's question. Select the most relevant and fact-based information:
        Knowledge Retrieval Responses:
        {responses}
        Pick only the best details to directly answer the customer's question clearly and accurately."""

        genai_prompt3_1 = """Example: If a customer asks, 'Tell me about green juice,' your response should be:
        'Our green juice, Well Greens, comes in a 6-pack for $41.94. It contains apple, spinach, kale, celery, lemon, and ginger juices, with 120 calories per serving.'"""

        genai_prompt3 = f"""{prompt}\n Only state the answer—do not include response numbers or extra commentary. 
        Be concise, factual, and avoid adding information beyond what is given. 
        If the answer is unavailable, say, 'I don't have that information at the moment.'"""
        

        new_prompt = genai_prompt+genai_prompt2+genai_prompt3_1+genai_prompt3

        return StreamingResponse(event_stream (model, new_prompt, ""), media_type="text/event-stream")
        
    except Exception as e:
        raise HTTPException(status_code=500, detail="Exception occurred: " + str(e))