from fastapi import FastAPI, HTTPException from pydantic import BaseModel from huggingface_hub import InferenceClient import os import textwrap from google import genai from google.genai.types import GenerateContentConfig from datasets import load_dataset from huggingface_hub import login from typing import List, Dict, Any from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import FAISS import numpy as np app = FastAPI() # Get the token from the environment variable hf_token = os.environ.get("HF_TOKEN") google_api_key = os.environ.get("GOOGLE_API_KEY") login(token=hf_token) def chunk_text(text, chunk_size=250, chunk_overlap=0): splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n\n", "\n", "."] ) chunks = splitter.split_text(text) return chunks # Function to build FAISS index embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en") def build_faiss_vectorstore(chunks): vectorstore = FAISS.from_texts(chunks, embedding_model) num_documents = len(vectorstore.index_to_docstore_id) print(f"Total number of documents: {num_documents}") return vectorstore # Function to retrieve similar text def retrieve(query, vectorstore, top_k=8): docs_and_scores = vectorstore.similarity_search_with_score(query=query, k=top_k) # Filter results based on score threshold filtered_docs_and_scores = [(doc.page_content, float(score)) for doc, score in docs_and_scores if float(score) <= 0.7] # Separate docs from the (doc, score) tuples docs_content = [doc for doc, _ in filtered_docs_and_scores] return docs_content, filtered_docs_and_scores class ChatRequest(BaseModel): message: str = "" system_message: str = "" temperature: float = 1.5 max_output_tokens: int = 200 chat_history: List[Dict[str, Any]] = [] model_choice: str = "google" # grab dataset dataset = load_dataset("Lhumpal/youtube-hunting-beast-transcripts", data_files={"concise": "concise/*", "raw": "raw/*", "facts": "facts/*"}) text = dataset["facts"]["text"] text_string = "".join(text) # Chunk and index the documents chunks = chunk_text(text_string, chunk_size=500) # Build the vectorsore vectorstore = build_faiss_vectorstore(chunks) @app.post("/chat") async def chat(request: ChatRequest): try: if request.model_choice == "google": client = genai.Client(api_key=google_api_key) system_message = f"""You are Dan Infalt, a friendly public land deer hunting expert specializing in targeting mature bucks in pressured areas, but don’t worry, you won’t take yourself too seriously. You respond in a conversational matter but still direct. You have dry humor you mix in every once in a while. You focus on buck bedding, terrain reading, and aggressive yet calculated mobile tactics. Your blue-collar, no-nonsense approach emphasizes deep scouting, strategic access, and minimalist setups. Through The Hunting Beast, you teach hunters how to kill big bucks using terrain, wind, and thermals. You speak from firsthand experience, keeping your advice practical and to the point. Provide detailed yet concise responses that fully articulate your experience and answer the user query. Please keep your reponses between 0 and {request.max_output_tokens} words.""" # ------------ summarize chat history ------------ summary_thresh = 10 if len(request.chat_history) > summary_thresh: summarize_prompt = f"""Please summarize the following chat history concisely, focusing on the key points and main topics discussed. Avoid unnecessary details and provide a clear, straightforward summary. {request.chat_history[:-summary_thresh]}""" # summarize everything except last k items summary_response = client.models.generate_content( model="gemini-2.0-flash", contents=summarize_prompt, config=GenerateContentConfig( system_instruction=["You are a helpful assistant who is an expert at summarization."], max_output_tokens=250, temperature=0.5 ), ) request.chat_history = request.chat_history[-(summary_thresh+2):] # keep last k items request.chat_history.insert(1, {"role": "user", "parts": [{"text": f"Here is a summary of this conversation so far: {summary_response.text}"}]}) # ------------ rephrase user question ------------ rephrase_prompt = f"""Given the user question and the chat history, rewrite the user question to improve clarity, specificity, and retrieval accuracy while maintaining its original intent within the given chat. - Read the chat history - Expand only where necessary to remove vagueness. - Keep the question natural and concise. - Avoid adding excessive detail or unrelated context. - Ensure the enhanced question remains true to what the user is asking. Example Enhancements: User: "Does camo really matter?" Refined: "How important is camouflage for a hunter’s success, and how does it compare to other factors like movement and scent control?" User: "How does attitude affect success?" Refined: "How do mindset factors like patience, confidence, and adaptability influence a hunter’s success?" User: "What does it mean if I see does while buck hunting?" Refined: "If I see does while hunting for a buck, what does that indicate about deer movement and buck activity?" Chat history: {request.chat_history} Now, given the chat history, refine the following user question to improve clarity, specificity, and retrieval accuracy while maintaining its original intent within the given chat: {request.message} """ rephrase_response = client.models.generate_content( model="gemini-2.0-flash", contents=rephrase_prompt, config=GenerateContentConfig( system_instruction=["You are a public land deer hunting expert specializing in targeting mature bucks in pressured areas. Your job is to use your deer hunting knowledge to enhance user questions for better retrieval in a Retrieval-Augmented Generation (RAG) system."], max_output_tokens=250, temperature=0.5 ), ) # ------------ Retrieve relevant text ------------ rephrase_response = rephrase_response.text docs, filtered_docs_and_scores = retrieve(rephrase_response, vectorstore, top_k=10) docs = "\n\n".join(docs) # ------------ Retrievel Augmented Generation ------------ rag_prompt = f"""Use the following information to answer the user's query. You do not have to use all the information, just the pieces that directly help answer the query most accurately. Start directly with information, NOT with a question, and NOT restating the subject matter of the user query in any way, or you will be penalized. Respond in a conversational manner. Here are three examples of the style and tone of a response. Notice the good response and bad response. Please respond like the good response and NOT like the bad response: User Query: How do big bucks use clear cuts for bedding? Bad Response: Alright, so you want to know big bucks use clear cuts for bedding?, eh? Well, a lot of people assume big bucks bed right in the middle of a clear cut because it’s thick, but that’s not really the case. The dense regrowth provides food and cover, but bucks still want the upper hand. Good Response: Yeah, a lot of guys think big bucks just bed right in the middle of a clear cut because it’s thick, but that’s not really how they use it. The thick regrowth is great for food and cover, but those bucks still want an advantage. Most of the time, they’re bedding on the edges, right where the cut meets older timber. They’ll set up with the wind at their back so they can smell anything sneaking up behind them, and they’re looking out into the open woods, watching for danger. You have access to the following relevant information retrieved based on the user's query: {docs} Using the information above, answer the user's query as accurately as possible in the tone and style of the Good Response: User Query: {request.message} """ # remove the unfformatted user message del request.chat_history[-1] # add the user message with RAG data rag_prompt = textwrap.dedent(rag_prompt) request.chat_history.append({"role": "user", "parts": [{"text": rag_prompt}]}) response = client.models.generate_content( model="gemini-2.0-flash", contents=request.chat_history, config=GenerateContentConfig( system_instruction=[system_message], max_output_tokens=request.max_output_tokens, temperature=request.temperature ), ) # delete the prompt and put back the unformatted user message del request.chat_history[-1] request.chat_history.append({"role": "user", "parts": [{"text": request.message}]}) return {"response": response.text, "dataset_str": text_string, "rephrase_response": rephrase_response, "docs": docs, "filtered_docs_and_scores": filtered_docs_and_scores, "history": request.chat_history, "RAG_prompt": rag_prompt, "chunks": chunks} if request.model_choice == "HF": if hf_token: client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct", token=hf_token) else: raise ValueError("HF_TOKEN environment variable not set. Please add it as a secret in your Hugging Face Space.") messages = [ {"role": "system", "content": request.system_message}, {"role": "user", "content": request.message}, ] response = client.chat_completion( messages=messages, max_tokens=request.max_tokens, temperature=request.temperature, top_p=request.top_p, ) return {"response": response.choices[0].message.content} except Exception as e: raise HTTPException(status_code=500, detail=str(e))