File size: 5,004 Bytes
0ab8e73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import os
import logging
import tempfile
from typing import List
import gradio as gr
# ---- LlamaIndex / Pinecone ----
from pinecone import Pinecone, ServerlessSpec
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext, Settings
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
# ---- Config ----
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# You can override these via Space "Variables" (Secrets)
PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "dds-demo-index")
PINECONE_REGION = os.getenv("PINECONE_REGION", "us-east-1") # keep in sync with ServerlessSpec
PINECONE_CLOUD = os.getenv("PINECONE_CLOUD", "aws")
EMBED_MODEL = os.getenv("EMBED_MODEL", "text-embedding-3-small")
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-4o-mini")
if not PINECONE_API_KEY:
raise RuntimeError("Missing PINECONE_API_KEY. Add it in your Space settings (Secrets).")
if not OPENAI_API_KEY:
raise RuntimeError("Missing OPENAI_API_KEY. Add it in your Space settings (Secrets).")
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("dds-space")
# ---- Pinecone client & index bootstrap ----
pc = Pinecone(api_key=PINECONE_API_KEY)
# Create index if it doesn't exist.
def _ensure_index(index_name: str, dimension: int = 1536):
existing = [idx["name"] for idx in pc.list_indexes()]
if index_name not in existing:
logger.info(f"Creating Pinecone index '{index_name}' (dim={dimension})...")
pc.create_index(
name=index_name,
dimension=dimension,
metric="cosine",
spec=ServerlessSpec(cloud=PINECONE_CLOUD, region=PINECONE_REGION),
)
return pc.Index(index_name)
pinecone_index = _ensure_index(PINECONE_INDEX_NAME, dimension=1536)
# ---- LlamaIndex settings ----
# Set global settings for LlamaIndex (embeddings + LLM)
Settings.embed_model = OpenAIEmbedding(model=EMBED_MODEL, api_key=OPENAI_API_KEY)
Settings.llm = OpenAI(model=LLM_MODEL, api_key=OPENAI_API_KEY)
# Vector store wrapper
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
def build_or_update_index(files: List[gr.File]) -> str:
"""
Load the uploaded files, chunk them with LlamaIndex, and upsert into Pinecone.
"""
if not files:
return "Please upload at least one file."
with tempfile.TemporaryDirectory() as tmpdir:
paths = []
for f in files:
# Gradio File object -> save to temp path
dst = os.path.join(tmpdir, os.path.basename(f.name))
with open(f.name, "rb") as src, open(dst, "wb") as out:
out.write(src.read())
paths.append(dst)
docs = SimpleDirectoryReader(input_files=paths).load_data()
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# Build a new index (will upsert into Pinecone via the vector_store)
_ = VectorStoreIndex.from_documents(
docs,
storage_context=storage_context,
show_progress=True,
)
return f"Indexed {len(files)} file(s) into Pinecone index: {PINECONE_INDEX_NAME}."
def answer(query: str, top_k: int = 4) -> str:
if not query or not query.strip():
return "Ask a question about your uploaded knowledge."
# Re-build a lightweight index wrapper that reads from the existing vector store
index = VectorStoreIndex.from_vector_store(vector_store)
qe = index.as_query_engine(similarity_top_k=top_k)
resp = qe.query(query)
return str(resp)
# ---- UI ----
INTRO = (
"Upload PDFs/TXT/Docs to build a Pinecone vector index (1536-d). "
"Then ask questions to retrieve & summarize with LlamaIndex + OpenAI."
)
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"<h1 style='text-align:center;'>📚 RAG with LlamaIndex + Pinecone</h1>"
"<p style='text-align:center;'>Omantel/DDS demo Space — minimal, production-friendly layout</p>"
)
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### 1) Upload & Index")
file_uploader = gr.File(label="Upload documents", file_count="multiple", type="filepath")
index_btn = gr.Button("Build / Update Index")
index_status = gr.Markdown()
with gr.Column(scale=1):
gr.Markdown("### 2) Ask a Question")
query = gr.Textbox(label="Your question", placeholder="e.g., What is the refund policy?")
topk = gr.Slider(1, 10, value=4, step=1, label="Top-K")
ask_btn = gr.Button("Ask")
answer_box = gr.Markdown()
gr.Markdown(f"**How it works:** {INTRO}")
index_btn.click(build_or_update_index, inputs=[file_uploader], outputs=[index_status])
ask_btn.click(answer, inputs=[query, topk], outputs=[answer_box])
if __name__ == "__main__":
demo.launch()
|