1 line
7.1 KiB
JSON
1 line
7.1 KiB
JSON
{"repo": "ashutoshpatel1/laptop-rag", "files": {"rag.py": {"5890c7bc8498cfe741b5794c0ef56f31af5f8724": {"t": 1785674360.2053196, "n": 6832, "h": "e3979f28f3187f35", "c": "import os\nimport google.generativeai as genai\nfrom dotenv import load_dotenv\nimport vector_store\n\n# Step 1: Load environment variables\n# We use load_dotenv() to read configuration options (like API keys) from the .env file.\nload_dotenv()\n\n# Step 2: Set the default API Key (supports Gemini and OpenRouter)\n# Gemini keys start with 'AIzaSy', while OpenRouter keys start with 'sk-or-v1-'\n# We default to the working OpenRouter key you provided.\nGEMINI_API_KEY = os.getenv(\"GEMINI_API_KEY\", \"\").strip()\n\n\n\ndef retrieve_top_k(query: str, k: int = 4) -> list:\n \n # Connect to the local ChromaDB database client\n client = vector_store.get_chroma_client()\n # Get our laptop specifications collection (table)\n collection = vector_store.get_collection(client)\n \n # Query ChromaDB using the question text.\n # ChromaDB will automatically convert the query text into an embedding (vector)\n # using our sentence-transformer model and look for the closest matches.\n results = collection.query(\n query_texts=[query],\n n_results=k\n )\n \n matched_docs = []\n # If we found any results, format them into a list of dictionaries for easier handling\n if results and results[\"documents\"]:\n docs = results[\"documents\"][0]\n metas = results[\"metadatas\"][0]\n distances = results[\"distances\"][0]\n for i in range(len(docs)):\n matched_docs.append({\n \"text\": docs[i],\n \"metadata\": metas[i],\n \"distance\": distances[i]\n })\n return matched_docs\n\n\ndef generate_answer(query: str, context_docs: list, api_key: str | None = None) -> str:\n \n \n context_text = \"\\n\\n\".join(\n [f\"Laptop Spec Details:\\n{doc['text']}\" for doc in context_docs]\n )\n \n \n system_prompt = (\n \"Instructions:\\n\"\n \"You are a friendly and helpful laptop sales assistant. \"\n \"Answer the user's question using the provided context of laptop specs. \"\n \"Find the best matching laptops in the context based on the user's request (e.g. price, brand, features). \"\n \"Even if the laptops in the context do not perfectly match the user's criteria (like budget or exact specs), \"\n \"recommend the closest alternatives from the context and explain how they compare to what they asked for. \"\n \"If the user's question is completely unrelated to laptops or tech, kindly explain that you can \"\n \"only assist with laptop specifications and comparisons. \"\n \"Do not invent any specifications, prices, or ratings that are not explicitly present in the context.\\n\\n\"\n )\n \n # 3. Create the final prompt with the system instructions, context, and query\n user_prompt = f\"{system_prompt}Context:\\n{context_text}\\n\\nQuestion: {query}\\n\\nAnswer:\"\n \n # 4. Resolve which API key to use (provided in UI or loaded from configuration/default)\n resolved_key = (api_key or GEMINI_API_KEY or \"\").strip()\n if not resolved_key:\n return \"Error: API Key is missing. Please set it in the sidebar or check your .env file.\"\n \n # 5. Route to OpenRouter if the key is in OpenRouter format\n if resolved_key.startswith(\"sk-or-v1-\"):\n import requests\n try:\n url = \"https://openrouter.ai/api/v1/chat/completions\"\n headers = {\n \"Authorization\": f\"Bearer {resolved_key}\",\n \"Content-Type\": \"application/json\"\n }\n payload = {\n \"model\": \"google/gemini-2.5-flash\",\n \"messages\": [\n {\"role\": \"user\", \"content\": user_prompt}\n ],\n \"temperature\": 0.0,\n \"max_tokens\": 1000 # Restrict tokens to prevent billing/credit check errors on free tier\n }\n response = requests.post(url, headers=headers, json=payload)\n if response.status_code == 200:\n result_json = response.json()\n answer = result_json[\"choices\"][0][\"message\"][\"content\"].strip()\n if not answer:\n return \"I don't know from the available info.\"\n return answer\n else:\n return f\"Error calling OpenRouter API: HTTP {response.status_code} - {response.text}\"\n except Exception as e:\n return f\"Error calling OpenRouter API: {e}\"\n \n # 6. Fallback to direct Gemini API using Google Generative AI package\n try:\n # Configure the Google GenAI library with our key\n genai.configure(api_key=resolved_key)\n # Load the Gemini model. We use gemini-2.0-flash as it is fast, cheap, and powerful.\n model = genai.GenerativeModel(model_name=\"gemini-2.0-flash\")\n \n # Call the API to generate content\n response = model.generate_content(\n user_prompt,\n # We set temperature to 0.0 to make the output highly predictable and factual\n generation_config={\"temperature\": 0.0}\n )\n \n answer = response.text.strip()\n if not answer:\n return \"I don't know from the available info.\"\n return answer\n except Exception as e:\n import traceback\n traceback.print_exc()\n return f\"Error calling Gemini API: {e}\"\n\n\ndef run_rag_pipeline(query: str, k: int = 4, api_key: str | None = None) -> dict:\n \"\"\"\n Step 5: Unified Pipeline\n Orchestrates the entire RAG flow:\n 1. Query ChromaDB to retrieve the most relevant laptops.\n 2. Run similarity distance check (guardrail) to reject queries with poor matches.\n 3. Call the LLM (Gemini or OpenRouter) to generate the factual response.\n \"\"\"\n # 1. Retrieve top-k documents from ChromaDB\n retrieved = retrieve_top_k(query, k=k)\n \n # 2. Guardrail check:\n # If the distance is too large (L2 distance > 1.35), it means the query is too different\n # from any laptop specs we have, so we refuse to answer to prevent hallucination.\n if not retrieved or retrieved[0][\"distance\"] > 1.35:\n return {\n \"answer\": \"I don't know from the available info.\",\n \"retrieved\": retrieved\n }\n \n # 3. Generate the final answer using the LLM (OpenRouter/Gemini)\n answer = generate_answer(query, retrieved, api_key=api_key)\n \n # Local fallback option if the API has connection/billing issues\n if \"Error calling Gemini\" in answer or \"Error calling OpenRouter\" in answer:\n top_laptop = retrieved[0]\n answer = (\n f\"[Local Fallback Answer - API Error] Best matched laptop: \"\n f\"{top_laptop['metadata'].get('brand')} {top_laptop['metadata'].get('model')} | \"\n f\"Specs: {top_laptop['text'][:200]}...\"\n )\n \n return {\n \"answer\": answer,\n \"retrieved\": retrieved\n }\n"}}}} |