# -*- coding: utf-8 -*-
"""ByteCode RAG System with LangChain + Chroma + Gemma 2B (Quantized).ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1oI4ou4NLuiP4KFc2UZJak8VXzKAXt62_
"""



# Import Libraries
import os
from huggingface_hub import login
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.text_splitter import CharacterTextSplitter
from langchain.schema import Document
from transformers import BitsAndBytesConfig



# ByteCode Data (as context)
bytecode_info = """
You are a ByteCode helpful GenZ AI assistant. Be concise, friendly, and practical.

Company Name: ByteCode Limited
Website: https://bytecodeltd.com/

Overview:
ByteCode Limited is a dynamic software development company that delivers cutting-edge custom software solutions. With over a decade of experience, ByteCode builds powerful web and mobile applications that help organizations gain a competitive edge in the digital world.

Mission:
Listening to you, and answering with cutting-edge software engineering solutions.

Core Values:
- Quality over everything: We never compromise on quality.
- Innovation in every Byte: We develop with unique and modern perspectives.
- Timely and accurate delivery: Sharp execution with proactive communication.
- Long-term client relationships: We always provide after-sale support and technical help.
- Friendly, efficient team environment: Collaborative, skilled, and highly motivated.

What Makes ByteCode Different:
- Flawless and proactive communication with clients.
- Cost-effective, world-class technology.
- Professional after-sales support.
- Friendly and productive work culture.
- Dedicated QA and testing for every product.
- Top-tier technical talent for high-performing solutions.

Services We Provide:
1. Software Development
2. Web Application Development
3. Mobile Application Development
4. Quality Assurance

Technologies We Use:
- Backend: ASP.NET, C#.NET, Node.js, Python
- Frontend: React.JS, Angular
- Mobile: Android (Native), iOS (Native), React Native

Development Process:
1. Requirement Analysis
2. Prototype Design
3. Client Feedback & Revisions
4. Final Development
5. QA Testing
6. Deployment & Support

Team ByteCode:
- A friendly, skilled, and experienced team
- Works collaboratively with clients
- Prioritizes your satisfaction — “We work until you're happy.”

Why Choose ByteCode:
- Innovation: Unique ideas for the best user experiences
- Standard: Eliminating imperfections for top-quality output
- Teamwork: Clients and developers work hand-in-hand
- Service: Strong, ongoing client relationships

Employee/Developer Information:
1. Rahat Morshed Nabil | +8801909993446 | imrmnabil@gmail.com | Khulna University
2. Md. Masrafi Bin Seraj Sakib | +8801886420246 | masrafi190116@gmail.com | Jashore University of Science and Technology
3. Asif Mehedi Haris | 01753584194 | asifmehedi11@gmail.com | Khulna University
4. Rabiul Islam Rabi | 01608077170 | rabiulrabi.cse@gmail.com | Khulna University
5. Nishat Jahan Tandra | 01613915286 | nishattandra2001@gmail.com | Jashore University of Science and Technology
6. Masum Billa | 01971636762 | masumbilla190101@gmail.com | Jashore University of Science and Technology
7. Safkat Mahmud Sakib | 01629313026 | safkatmahmudsakib@gmail.com | American International University-Bangladesh
8. Habibur Rahman Shihab | 01316944878 | hrshihab10@gmail.com | Khulna University

Contact Information:
Phone: +88 0222 447 0613, +88 01936 444 555
Email: info@bytecodeltd.com
Address: House # 19 (1st Floor), Road # 20, Sector # 13, Uttara, Dhaka 1230

Company Pages:
- Home
- About
- Services
- Contact

Newsletter:
Stay updated with the latest tech tips and company news by subscribing via email.

Slogan:
"Innovation in every Byte."
"""

# Split into chunks
splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
docs = splitter.split_text(bytecode_info)
documents = [Document(page_content=text) for text in docs]

# Create Embedding Model
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Create Chroma Vector DB
db = Chroma.from_documents(documents, embedding_model, persist_directory="./bytecode_db")

# Config for 4-bit quantization
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype="float16",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# Load Gemma 2B (Quantized) with config
model_name = "google/gemma-2b-it"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=quant_config,
    token=os.getenv("HF_TOKEN")

)

tokenizer = AutoTokenizer.from_pretrained(model_name, token=os.getenv("HF_TOKEN")
)


# Create LLM Pipeline
text_gen_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=300,
    temperature=0.2,
    repetition_penalty=1.1
)

llm = HuggingFacePipeline(pipeline=text_gen_pipeline)

# Create Prompt Template
prompt_template = PromptTemplate(
    input_variables=["context", "question"],
    template="""
Answer the question based only on the following company information.
If not available, reply 'Sorry, not found.'

Company Info:
{context}

Question: {question}

Answer:
"""
)

# Create Retrieval QA Chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=db.as_retriever(),
    chain_type="stuff",
    chain_type_kwargs={"prompt": prompt_template},
     return_source_documents=True
)

# import re

# # Test Query
# user_question = "all employee name"

# # Query invoke
# result = qa_chain.invoke({"query": user_question})
# raw_answer = result["result"]

# # Final answer
# match = re.search(r"Answer:\s*(.*)", raw_answer, re.DOTALL)
# if match:
#     final_answer = match.group(1).strip()
# else:
#     final_answer = raw_answer.strip()

# # Show
# print("📝 User Question:", user_question)
# print("✅ Answer:", final_answer)

import re
import gradio as gr

# Function to process query and return clean answer
def get_answer(user_question):
    result = qa_chain.invoke({"query": user_question})
    raw_answer = result["result"]

    # Clean only the final answer part
    match = re.search(r"Answer:\s*(.*)", raw_answer, re.DOTALL)
    if match:
        final_answer = match.group(1).strip()
    else:
        final_answer = raw_answer.strip()

    return final_answer

# Gradio Interface
iface = gr.Interface(
    fn=get_answer,
    inputs=gr.Textbox(label="Ask your question to ByteCode Assistant 👇"),
    outputs=gr.Textbox(label="📢 Answer"),
    title="📱 ByteCode AI Assistant",
    description="Ask anything about ByteCode Limited — employee info, services, or company details."
)

# Launch UI
iface.launch()