import chromadb
from chromadb.utils import embedding_functions

# Initialize Chroma client
client = chromadb.PersistentClient(path="./chroma_db")

# Create embedding function
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# Create collection
collection = client.get_or_create_collection(
    name="support_docs",
    embedding_function=embedding_fn
)

# Add documents
documents = [
    "How to reset password...",
    "What are your shipping policies?...",
    "How to track an order?..."
]

ids = [f"doc_{i}" for i in range(len(documents))]

collection.add(
    documents=documents,
    ids=ids
)

print(f"Added {len(documents)} documents to vector database")

# Query the vector database
def search_docs(query, n_results=3):
    results = collection.query(
        query_texts=[query],
        n_results=n_results
    )
    return results['documents'][0]

# Test search
query = "How do I reset my password?"
relevant_docs = search_docs(query)
print("Relevant documents:")
for doc in relevant_docs:
    print(f"- {doc}")

from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA

# Setup
embeddings = OpenAIEmbeddings()
llm = ChatOpenAI(model="gpt-3.5-turbo")

# Load vector store
vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings
)

# Create RAG chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3})
)

# Query
query = "How do I reset my password?"
answer = qa_chain.run(query)
print(answer)

# Custom RAG without LangChain
import openai

def rag_answer(query):
    # 1. Retrieve relevant docs
    relevant_docs = search_docs(query)
    context = "\n".join(relevant_docs)

    # 2. Create prompt
    prompt = f"""You are a customer support assistant.
    Use the following context to answer the question.
    If the answer isn't in the context, say you don't know.

    Context:
    {context}

    Question: {query}
    Answer:"""

    # 3. Generate answer
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# Test
answer = rag_answer("How do I reset my password?")
print(answer)

import streamlit as st

st.title("Customer Support Assistant")
st.markdown("Ask questions about our products and services.")

# Initialize RAG system
@st.cache_resource
def load_rag():
    # Load vector store and RAG chain
    return qa_chain

qa_chain = load_rag()

# User input
query = st.text_input("Your question:", placeholder="e.g., How do I reset my password?")

if query:
    with st.spinner("Searching for answers..."):
        answer = qa_chain.run(query)

    st.subheader("Answer")
    st.write(answer)

    # Show sources
    with st.expander("Sources"):
        docs = vectorstore.similarity_search(query, k=3)
        for doc in docs:
            st.write(f"- {doc.page_content[:200]}...")

# Deploy on Streamlit Cloud
# 1. Create requirements.txt with:
#    streamlit
#    chromadb
#    langchain
#    openai

# 2. Add OPENAI_API_KEY to secrets
# 3. Push to GitHub and deploy

# Environment variables (secrets.toml)
# OPENAI_API_KEY = "sk-..."

# Run locally:
# streamlit run app.py