app.py

# Import the necessary Libraries 
import os 
import uuid
import json
import traceback
import gradio as gr
import requests
from openai import OpenAI
from chromadb import PersistentClient
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma

from huggingface_hub import CommitScheduler
from pathlib import Path
from dotenv import load_dotenv


# Use a simpler and portable environment variable name
#secret_key_name = "ACCOUNTS_P3_SECRET"


# Load environment variables
load_dotenv()

# Get API key from env
api_key = os.getenv("FIREWORKS_API_KEY")
if not api_key:
    raise EnvironmentError("Missing FIREWORKS_API_KEY")

# Create OpenAI client for Fireworks
client = OpenAI(
    base_url="https://api.fireworks.ai/inference/v1",
    api_key=api_key,
)
    
url = "https://api.fireworks.ai/inference/v1/chat/completions"
payload = {
   "model": "accounts/fireworks/models/llama-v3p3-70b-instruct",
    "max_tokens": 1024,
    "temperature": 0.6,
    "top_p": 1,
    "top_k": 40,
    "presence_penalty": 0,
    "frequency_penalty": 0,
    "messages": [
        {
            "role": "user",
            "content": "Hello, how are you?"
        }
    ]
}
headers = {
    "Accept": "application/json",
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
}

response = requests.post(url, headers=headers, json=payload)

if response.ok:
    print(response.json()["choices"][0]["message"]["content"])
else:
    print(f"❌ Error {response.status_code}: {response.text}")

# # Make a test call
# response = client.chat.completions.create(
#     model="accounts/fireworks/models/llama-v3p1-405b-instruct",
#     messages=[{"role": "user", "content": "Hello, Fireworks!"}]
# )

#print(response.choices[0].message.content)


embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
# Define the embedding model and the vectorstore

# client = PersistentClient(path="./report_10kdb")
# print("Available collections:", client.list_collections())


# collection_name = 'report-10k-2024'
collection_name ='report_10k_collection'
vectorstore_persisted = Chroma(
    collection_name=collection_name,
    persist_directory='./report_10kdb',
    embedding_function=embedding_model
)


# Load the persisted vectorDB

retriever = vectorstore_persisted.as_retriever(
    search_type='similarity',
    search_kwargs={'k': 5}
)


# Prepare the logging functionality

log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
log_folder = log_file.parent

scheduler = CommitScheduler(
    repo_id="RAG-investment-recommendation-log",
    repo_type="dataset",
    folder_path=log_folder,
    path_in_repo="data",
    every=2
)

# Define the Q&A system message

qna_system_message = """
You are an assistant to a researcher. Your task is to provide relevant information about The 10K reports repository.

User input will include the necessary context for you to answer their questions. This context will begin with the token: ###Context.
The context contains references to specific portions of documents relevant to the user's query, along with source links.
The source for a context will begin with the token ###Source

When crafting your response:
1. Select the most relevant context or contexts to answer the question.
2. Include the source links in your response, document name and page number.
3. User questions will begin with the token: ###Question.
4. If the question is irrelevant to 10k report respond with - "Apologies, I can only help you with questions related to the 10k Reports."

Please adhere to the following guidelines:
- Answer only using the context provided. If you do not know the answer say 'Sorry,I do not know.'
- Do not mention anything about the context in your final answer.
- If the answer is not found in the context, it is very important for you to respond with "I don't know. Please check the docs found in the report repository."
- Always quote the source when you use the context. Cite the relevant source at the end of your response under the section - Sources:
- Do not make up sources. Use the links provided in the sources section of the context and nothing else. You are prohibited from providing other links/sources or general knowledge about people not in the 10k reports.

Here is an example of how to structure your response:

Answer:
[Answer]

Source
[Source]
"""

# Define the user message template
qna_user_message_template = qna_user_message_template = """
###Question:
{question}

###Context:
{context}

###Instructions:
Please answer the question using only the information in the context above.
- If the answer is not in the context, say: "I don't know. Please check the docs found in the report repository."
- Always include a 'Sources:' section listing relevant source links from the context.
- Do not add external knowledge or invent sources.
- Respond in the format below:

###Question:
{question}

###Context:
{context}



Answer:
[Your answer here]

###Context:
{context}

Sources:
[List of sources used]
"""

# Define the predict function that runs when 'Submit' is clicked or when a API request is made
def predict(user_input,company):
    # print(f'{vectorstore_persisted.get(include=["metadatas", "documents"])}')   # new
    api_key = os.getenv("FIREWORKS_API_KEY")
    filter_text = "dataset/"+company+"-10-k-2023.pdf"
    
   # relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter_text})
   
    # Create context_for_query
    relevant_document_chunks = retriever.get_relevant_documents(user_input) #New
    context_list = [d.page_content for d in relevant_document_chunks]
    context_for_query = ". ".join(context_list)

    # Create messages
    # print(f"Context: {context_for_query[:5000]}\n end")
    prompt = [
        {'role':'system', 'content': qna_system_message},
        {'role': 'user', 'content': qna_user_message_template.format(
            context=context_for_query,
            question=user_input
            )
        }
    ]
    
    # Get response from the LLM
    
    # Get response from the LLM
    try:
        response = client.chat.completions.create(
            model= "accounts/fireworks/models/llama-v3p3-70b-instruct",
            messages=prompt,
            temperature=0
        )

        prediction = response.choices[0].message.content

    except Exception as e:
        prediction = {
            "error_type": type(e).__name__,
            "error_message": str(e),
            "traceback": traceback.format_exc()
        }
        
    # # Log safely
    # with scheduler.lock:
    #     with log_file.open("a") as f:
    #         f.write(json.dumps(
    #             {
    #                 'user_input': user_input,
    #                 'retrieved_context': context_for_query,
    #                 'model_response': prediction
    #             },
    #             indent=2
    #         ))
    #         f.write("\n")
    
    return prediction

examples = [
    ["What are the risks factors Amazon faces from The People’s Republic of China?", "AWS"],
    ["What are the primary business segments of AWS, and how does each segment contribute to the overall revenue and profitability?", "AWS"],
    ["What are the key risk factors identified in the 10-K report that could potentially impact the AWS business operations and financial performance?", "AWS"],
    ["Has the company made any significant acquisitions in the AI space, and how are these acquisitions being integrated into the company's strategy?", "Microsoft"],
    ["How much capital has been allocated towards AI research and development?","Google"],
    ["What initiatives has IBM implemented to address ethical concerns surrounding AI, such as fairness, accountability, and privacy?","IBM"],
    ["How does Meta plan to differentiate itself in the AI space relative to competitors?","Meta"]
]

def get_predict(question, company):
    # Implement your prediction logic here
    if company == "AWS":
        # Perform prediction for AWS
        selectedCompany = "aws"
    elif company == "IBM":
        # Perform prediction for IBM
        selectedCompany = "IBM"
    elif company == "Google":
        # Perform prediction for Google
       selectedCompany = "Google"
    elif company == "Meta":
        # Perform prediction for Meta
        selectedCompany = "meta"
    elif company == "Microsoft":
        # Perform prediction for Microsoft
        selectedCompany = "msft"
    else:
        return "Invalid company selected"

    output = predict(question, selectedCompany)    
    return output

# Set-up the Gradio UI
# Add text box and radio button to the interface
# The radio button is used to select the company 10k report in which the context needs to be retrieved.

# Create the interface
# For the inputs parameter of Interface provide [textbox,company]

with gr.Blocks(theme="Taithrah/Minimal@>=0.0.1,<0.1.0") as demo:
    with gr.Row():
        company = gr.Radio(["AWS", "IBM", "Google", "Meta", "Microsoft"], label="Select a company")
    with gr.Row():
        question = gr.Textbox(label="Enter your question")

    submit = gr.Button("Submit")
    output = gr.Textbox(label="Output")

    submit.click(
        fn=get_predict,
        inputs=[question, company],
        outputs=output
    )
    examples_component = gr.Examples(examples=examples, inputs=[question, company]) 
demo.queue()
demo.launch()