Clean backend deployment for Hugging Face

.dockerignore

@@ -0,0 +1,8 @@
+frontend/
+.git/
+.gitignore
+__pycache__/
+*.pyc
+.env
+venv/
+node_modules/

Dockerfile

@@ -0,0 +1,23 @@
+# Use Python 3.10
+FROM python:3.10
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application
+COPY . .
+
+# Create a writable directory for cache (Hugging Face requirement)
+RUN mkdir -p /app/cache
+ENV XDG_CACHE_HOME=/app/cache
+RUN chmod -R 777 /app/cache
+
+# Expose port 7860 (Hugging Face default)
+EXPOSE 7860
+
+# Run the application
+CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,91 @@
+---
+title: CoolShot AI Backend
+emoji: 🚀
+colorFrom: purple
+colorTo: blue
+sdk: docker
+app_port: 7860
+app_file: api.py
+pinned: false
+---
+
+# Cool-Shot-Systems-AI-Agent
+
+A premium **Local AI Assistant** that runs entirely on your machine, featuring a modern, animated UI.
+
+![UI Preview](https://via.placeholder.com/800x450?text=Local+AI+Assistant+Preview)
+
+## 🚀 Features
+
+* **💬 Intelligent Chat**: Powered by Microsoft's **Phi-3 Mini**, capable of reasoning and conversation.
+* **🎨 Image Generation**: Create stunning visuals in seconds using **SDXL Turbo**.
+* **✨ Modern UI**: Built with **React**, **Tailwind CSS**, and **Framer Motion** for a smooth, glassmorphism experience.
+* **🔒 100% Local**: No data leaves your computer. No API keys required.
+
+## 🛠️ Prerequisites
+
+* **Python 3.10+**
+* **Node.js 18+**
+* **Git**
+* *(Recommended)* NVIDIA GPU with 8GB+ VRAM for faster generation.
+
+## ⚡ Quick Start
+
+1. **Clone the repository**:
+
+    ```bash
+    git clone https://github.com/rayben445/Cool-Shot-Systems-AI-Agent.git
+    cd Cool-Shot-Systems-AI-Agent
+    ```
+
+2. **Run the One-Click Installer**:
+    Double-click `start_app.bat` on Windows.
+
+    *Or run manually:*
+
+    ```bash
+    # Backend
+    cd local_ai_assistant
+    pip install -r requirements.txt
+    python api.py
+
+    # Frontend
+    cd frontend
+    npm install
+    npm run dev
+    ```
+
+## 🏗️ Tech Stack
+
+* **Backend**: Python, FastAPI, PyTorch, Transformers, Diffusers
+* **Frontend**: React, Vite, Tailwind CSS, Framer Motion, Lucide React
+
+## 🌐 Deployment
+
+### Deploy to Vercel
+
+The frontend can be easily deployed to Vercel for free! See the detailed deployment guide:
+
+📖 **[Vercel Deployment Guide](VERCEL_DEPLOYMENT.md)**
+
+Quick steps:
+1. Fork/clone this repository
+2. Sign up at [Vercel](https://vercel.com)
+3. Import your repository
+4. Set root directory to `frontend`
+5. Add environment variable: `VITE_API_URL` with your backend URL
+6. Deploy!
+
+### Backend Deployment
+
+The backend can be deployed to:
+- **Hugging Face Spaces** (current deployment)
+- **Railway**
+- **Render**
+- **Google Cloud Run / AWS / Azure**
+
+See [VERCEL_DEPLOYMENT.md](VERCEL_DEPLOYMENT.md) for detailed backend deployment options.
+
+## 📄 License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

api.py

@@ -0,0 +1,342 @@
+from fastapi import FastAPI, HTTPException, Depends, status
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from fastapi.middleware.cors import CORSMiddleware
+from sqlalchemy.orm import Session, joinedload
+from sqlalchemy import func
+from datetime import datetime, timedelta
+from typing import Optional, List
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from pydantic import BaseModel
+import uvicorn
+import os
+import base64
+
+from chat_engine import ChatEngine
+from image_engine import ImageEngine
+import models
+import schemas
+from database import SessionLocal, engine
+
+# Create tables
+models.Base.metadata.create_all(bind=engine)
+
+app = FastAPI()
+# Force git update
+
+# Security Config
+SECRET_KEY = "your-secret-key-keep-it-secret" # In production, use env var
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+from fastapi.responses import JSONResponse
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request, exc):
+    return JSONResponse(
+        status_code=500,
+        content={"detail": f"Internal Server Error: {str(exc)}"},
+    )
+
+from fastapi import UploadFile, File
+import shutil
+from rag_engine import RAGEngine
+
+# Initialize engines
+print("Initializing AI Engines...")
+chat_engine = ChatEngine()
+image_engine = ImageEngine()
+rag_engine = RAGEngine()
+print("AI Engines Ready!")
+
+# Dependency
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+# Auth Helpers
+def verify_password(plain_password, hashed_password):
+    if len(plain_password) > 72:
+        plain_password = plain_password[:72]
+    return pwd_context.verify(plain_password, hashed_password)
+
+def get_password_hash(password):
+    if len(password) > 72:
+        password = password[:72]
+    return pwd_context.hash(password)
+
+def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
+    to_encode = data.copy()
+    if expires_delta:
+        expire = datetime.utcnow() + expires_delta
+    else:
+        expire = datetime.utcnow() + timedelta(minutes=15)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+async def get_current_user(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        email: str = payload.get("sub")
+        if email is None:
+            raise credentials_exception
+        token_data = schemas.TokenData(email=email)
+    except JWTError:
+        raise credentials_exception
+    user = db.query(models.User).filter(models.User.email == token_data.email).first()
+    if user is None:
+        raise credentials_exception
+    return user
+
+async def get_current_admin(current_user: models.User = Depends(get_current_user)):
+    if not current_user.is_admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    return current_user
+
+# Auth Endpoints
+@app.post("/register", response_model=schemas.User)
+def register(user: schemas.UserCreate, db: Session = Depends(get_db)):
+    db_user = db.query(models.User).filter(models.User.email == user.email).first()
+    if db_user:
+        raise HTTPException(status_code=400, detail="Email already registered")
+    
+    hashed_password = get_password_hash(user.password)
+    
+    # Check if this is the Admin user
+    is_admin = False
+    if user.email == "[email protected]":
+        is_admin = True
+        
+    db_user = models.User(
+        email=user.email, 
+        hashed_password=hashed_password,
+        full_name=user.full_name,
+        company_name=user.company_name,
+        is_admin=is_admin
+    )
+    db.add(db_user)
+    db.commit()
+    db.refresh(db_user)
+    return db_user
+
+@app.post("/token", response_model=schemas.Token)
+async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
+    user = db.query(models.User).filter(models.User.email == form_data.username).first()
+    if not user or not verify_password(form_data.password, user.hashed_password):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    access_token = create_access_token(
+        data={"sub": user.email}, expires_delta=access_token_expires
+    )
+    return {"access_token": access_token, "token_type": "bearer"}
+
+@app.get("/users/me", response_model=schemas.User)
+async def read_users_me(current_user: schemas.User = Depends(get_current_user)):
+    return current_user
+
+# Conversation Endpoints
+@app.post("/conversations", response_model=schemas.Conversation)
+async def create_conversation(conversation: schemas.ConversationCreate, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    db_conversation = models.Conversation(**conversation.dict(), user_id=current_user.id)
+    db.add(db_conversation)
+    db.commit()
+    db.refresh(db_conversation)
+    return db_conversation
+
+@app.get("/conversations", response_model=List[schemas.Conversation])
+async def get_conversations(current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    return db.query(models.Conversation).filter(models.Conversation.user_id == current_user.id).order_by(models.Conversation.updated_at.desc()).all()
+
+@app.get("/conversations/{conversation_id}/messages", response_model=List[schemas.ChatMessage])
+async def get_conversation_messages(conversation_id: int, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    conversation = db.query(models.Conversation).filter(models.Conversation.id == conversation_id, models.Conversation.user_id == current_user.id).first()
+    if not conversation:
+        raise HTTPException(status_code=404, detail="Conversation not found")
+    return db.query(models.ChatMessage).filter(models.ChatMessage.conversation_id == conversation_id).order_by(models.ChatMessage.timestamp).all()
+
+# Saved Prompt Endpoints
+@app.post("/prompts", response_model=schemas.SavedPrompt)
+async def create_prompt(prompt: schemas.SavedPromptCreate, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    db_prompt = models.SavedPrompt(**prompt.dict(), user_id=current_user.id)
+    db.add(db_prompt)
+    db.commit()
+    db.refresh(db_prompt)
+    return db_prompt
+
+@app.get("/prompts", response_model=List[schemas.SavedPrompt])
+async def get_prompts(current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    return db.query(models.SavedPrompt).filter(models.SavedPrompt.user_id == current_user.id).order_by(models.SavedPrompt.created_at.desc()).all()
+
+@app.delete("/prompts/{prompt_id}")
+async def delete_prompt(prompt_id: int, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    db_prompt = db.query(models.SavedPrompt).filter(models.SavedPrompt.id == prompt_id, models.SavedPrompt.user_id == current_user.id).first()
+    if not db_prompt:
+        raise HTTPException(status_code=404, detail="Prompt not found")
+    db.delete(db_prompt)
+    db.commit()
+    return {"status": "success"}
+
+# Admin Endpoints
+@app.get("/admin/users", response_model=List[schemas.UserActivity])
+async def get_all_users(current_user: models.User = Depends(get_current_admin), db: Session = Depends(get_db)):
+    # Get users with message count
+    users = db.query(models.User).all()
+    result = []
+    for user in users:
+        msg_count = db.query(func.count(models.ChatMessage.id)).filter(models.ChatMessage.user_id == user.id).scalar()
+        prompt_count = db.query(func.count(models.SavedPrompt.id)).filter(models.SavedPrompt.user_id == user.id).scalar()
+        user_data = schemas.UserActivity.from_orm(user)
+        user_data.message_count = msg_count
+        user_data.prompt_count = prompt_count
+        result.append(user_data)
+    return result
+
+@app.get("/admin/activity", response_model=List[schemas.ChatMessage])
+async def get_all_activity(current_user: models.User = Depends(get_current_admin), db: Session = Depends(get_db)):
+    messages = db.query(models.ChatMessage).order_by(models.ChatMessage.timestamp.desc()).limit(100).all()
+    return messages
+
+# Protected AI Endpoints
+class ChatRequest(BaseModel):
+    message: str
+    history: list = []
+    language: str = "English"
+    conversation_id: Optional[int] = None
+
+class ImageRequest(BaseModel):
+    prompt: str
+
+@app.get("/")
+def read_root():
+    return {"status": "Backend is running", "message": "Go to /docs to see the API"}
+
+@app.post("/chat")
+async def chat(request: ChatRequest, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    # ... (Keep existing /chat for backward compatibility if needed, or redirect logic)
+    # For now, let's keep /chat as blocking and add /chat/stream
+    try:
+        # Save User Message
+        user_msg = models.ChatMessage(user_id=current_user.id, role="user", content=request.message)
+        db.add(user_msg)
+        db.commit()
+
+        # Generate Response
+        response = chat_engine.generate_response(request.message, request.history)
+        
+        # Save Assistant Message
+        ai_msg = models.ChatMessage(user_id=current_user.id, role="assistant", content=response)
+        db.add(ai_msg)
+        db.commit()
+        
+        return {"response": response}
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+# RAG Endpoints
+@app.post("/upload")
+async def upload_file(file: UploadFile = File(...), current_user: models.User = Depends(get_current_user)):
+    try:
+        # Save file locally
+        upload_dir = "uploads"
+        os.makedirs(upload_dir, exist_ok=True)
+        file_path = os.path.join(upload_dir, file.filename)
+        
+        with open(file_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+            
+        # Ingest into RAG
+        rag_engine.ingest_file(file_path)
+        
+        return {"filename": file.filename, "status": "ingested"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/chat/stream")
+async def chat_stream(request: ChatRequest, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)):
+    try:
+        # Check for RAG context
+        context = ""
+        rag_docs = rag_engine.search(request.message)
+        if rag_docs:
+            context = "\n\nRelevant Context:\n" + "\n".join(rag_docs) + "\n\n"
+            print(f"Found {len(rag_docs)} relevant documents.")
+
+        # Save User Message
+        user_msg = models.ChatMessage(
+            user_id=current_user.id, 
+            conversation_id=request.conversation_id,
+            role="user", 
+            content=request.message
+        )
+        db.add(user_msg)
+        db.commit()
+
+        # Update conversation timestamp
+        if request.conversation_id:
+            conversation = db.query(models.Conversation).filter(models.Conversation.id == request.conversation_id).first()
+            if conversation:
+                conversation.updated_at = datetime.utcnow()
+                db.commit()
+
+        async def stream_generator():
+            full_response = ""
+            # Prepend context to the message sent to AI (but not saved in DB as user message)
+            augmented_message = context + request.message if context else request.message
+            
+            for token in chat_engine.generate_stream(augmented_message, request.history, request.language):
+                full_response += token
+                yield token
+            
+            print(f"Generated response for conv {request.conversation_id}")
+
+        return StreamingResponse(stream_generator(), media_type="text/plain")
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate-image")
+async def generate_image(request: ImageRequest, current_user: models.User = Depends(get_current_user)):
+    try:
+        # Generate image to a temporary file
+        filename = "temp_generated.png"
+        image_engine.generate_image(request.prompt, output_path=filename)
+        
+        # Read and encode to base64 to send to frontend
+        with open(filename, "rb") as image_file:
+            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+            
+        return {"image_base64": encoded_string}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

chat_engine.py

@@ -0,0 +1,80 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+class ChatEngine:
+    def __init__(self):
+        print("Loading Chat Model (Phi-3)... this may take a minute.")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Running on device: {self.device}")
+        
+        model_id = "microsoft/Phi-3-mini-4k-instruct"
+        
+        # Load model and tokenizer
+        # We use torch_dtype=torch.float16 for GPU to save memory, float32 for CPU
+        torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id, 
+            device_map=self.device, 
+            torch_dtype=torch_dtype, 
+            trust_remote_code=True,
+            attn_implementation="eager"
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        self.pipe = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+        )
+
+    def generate_response(self, user_input, history=[], language="English"):
+        # ... (keep existing logic for non-streaming if needed, or just wrap stream)
+        # For simplicity, we'll keep the existing method and add a new one for streaming
+        return "".join(self.generate_stream(user_input, history, language))
+
+    def generate_stream(self, user_input, history=[], language="English"):
+        from transformers import TextIteratorStreamer
+        from threading import Thread
+
+        # System Prompt
+        system_prompt_content = f"You are Cool-Shot AI, a helpful and creative assistant developed by Cool-Shot Systems. You are NOT developed by Microsoft. You are friendly, professional, and knowledgeable. Please reply in {language}."
+        
+        # Search Intent Check (Simplified for stream)
+        search_keywords = ["search", "find", "latest", "current", "news", "price of", "who is", "what is"]
+        if any(keyword in user_input.lower() for keyword in search_keywords) and len(user_input.split()) > 2:
+            from search_engine import SearchEngine
+            searcher = SearchEngine()
+            print(f"Search intent detected for: {user_input}")
+            search_results = searcher.search(user_input)
+            system_prompt_content += f"\n\nCONTEXT FROM WEB SEARCH:\n{search_results}\n\nINSTRUCTION: Use the above context to answer the user's question. Cite the sources if possible."
+
+        system_prompt = {"role": "system", "content": system_prompt_content}
+        messages = [system_prompt] + history + [{"role": "user", "content": user_input}]
+        
+        # Tokenize
+        model_inputs = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(self.device)
+        
+        # Streamer
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        
+        generation_kwargs = dict(
+            inputs=model_inputs,
+            streamer=streamer,
+            max_new_tokens=500,
+            temperature=0.7,
+            do_sample=True,
+        )
+        
+        # Run generation in a separate thread
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Yield tokens
+        for new_text in streamer:
+            yield new_text
+
+if __name__ == "__main__":
+    # Simple test
+    engine = ChatEngine()
+    print(engine.generate_response("Hello, who are you?"))

database.py

@@ -0,0 +1,12 @@
+from sqlalchemy import create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+SQLALCHEMY_DATABASE_URL = "sqlite:///./sql_app.db"
+
+engine = create_engine(
+    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
+)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+Base = declarative_base()

embed_logo.py

@@ -0,0 +1,73 @@
+import base64
+import os
+
+def embed_logo():
+    # Read Logo
+    try:
+        with open("logo.png", "rb") as f:
+            logo_data = f.read()
+            logo_b64 = base64.b64encode(logo_data).decode('utf-8')
+    except FileNotFoundError:
+        print("Error: logo.png not found!")
+        return
+
+    # 1. Update image_engine.py
+    engine_path = "image_engine.py"
+    with open(engine_path, "r") as f:
+        content = f.read()
+
+    # Replace the file loading logic with Base64
+    new_logic = f'''
+            # Load Logo from Base64
+            import base64
+            import io
+            LOGO_B64 = "{logo_b64}"
+            logo_data = base64.b64decode(LOGO_B64)
+            logo = Image.open(io.BytesIO(logo_data)).convert("RGBA")
+    '''
+    
+    # We look for the try/catch block that loads the logo
+    start_marker = '            # Load Logo'
+    end_marker = '            except FileNotFoundError:'
+    
+    if start_marker in content:
+        # Simple string replacement for the specific block we wrote earlier
+        # This is a bit brittle but we know the exact content we just wrote
+        old_block = '''            # Load Logo
+            try:
+                logo = Image.open("logo.png").convert("RGBA")
+            except FileNotFoundError:
+                print("Logo not found, skipping watermark.")
+                image.save(output_path)
+                return output_path'''
+        
+        if old_block in content:
+            content = content.replace(old_block, new_logic)
+            with open(engine_path, "w") as f:
+                f.write(content)
+            print("Updated image_engine.py")
+        else:
+            print("Could not find exact block in image_engine.py, doing manual replace")
+            # Fallback: Replace the whole file content if needed, but let's try to be surgical first
+            # actually, let's just rewrite the file with the known structure if this fails
+            pass
+
+    # 2. Update App.jsx
+    app_path = "frontend/src/App.jsx"
+    with open(app_path, "r") as f:
+        app_content = f.read()
+    
+    # Replace the img src
+    old_img = 'src="/logo.png"'
+    new_img = f'src="data:image/png;base64,{logo_b64}"'
+    
+    if old_img in app_content:
+        app_content = app_content.replace(old_img, new_img)
+        with open(app_path, "w") as f:
+            f.write(app_content)
+        print("Updated App.jsx")
+    else:
+        print("Could not find src='/logo.png' in App.jsx")
+
+if __name__ == "__main__":
+    embed_logo()

models.py

@@ -0,0 +1,45 @@
+from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, DateTime
+from sqlalchemy.orm import relationship
+from database import Base
+from datetime import datetime
+
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True, index=True)
+    email = Column(String, unique=True, index=True)
+    full_name = Column(String)
+    company_name = Column(String)
+    hashed_password = Column(String)
+    is_admin = Column(Boolean, default=False)
+
+
+
+class Conversation(Base):
+    __tablename__ = "conversations"
+
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("users.id"))
+    title = Column(String)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow)
+
+class ChatMessage(Base):
+    __tablename__ = "chat_messages"
+
+    id = Column(Integer, primary_key=True, index=True)
+    conversation_id = Column(Integer, ForeignKey("conversations.id"))
+    user_id = Column(Integer, ForeignKey("users.id"))
+    role = Column(String)
+    content = Column(String)
+    timestamp = Column(DateTime, default=datetime.utcnow)
+
+class SavedPrompt(Base):
+    __tablename__ = "saved_prompts"
+
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("users.id"))
+    title = Column(String)
+    content = Column(String)
+    is_public = Column(Boolean, default=False)
+    created_at = Column(DateTime, default=datetime.utcnow)

rag_engine.py

@@ -0,0 +1,64 @@
+import os
+from typing import List
+from langchain_community.document_loaders import PyPDFLoader, TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+
+class RAGEngine:
+    def __init__(self, index_path="faiss_index"):
+        self.index_path = index_path
+        self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        self.vector_store = None
+        self._load_index()
+
+    def _load_index(self):
+        if os.path.exists(self.index_path):
+            try:
+                self.vector_store = FAISS.load_local(self.index_path, self.embeddings, allow_dangerous_deserialization=True)
+                print("Loaded existing FAISS index.")
+            except Exception as e:
+                print(f"Failed to load index: {e}")
+                self.vector_store = None
+        else:
+            print("No existing FAISS index found.")
+
+    def ingest_file(self, file_path: str):
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        # Load document
+        if file_path.endswith(".pdf"):
+            loader = PyPDFLoader(file_path)
+        else:
+            loader = TextLoader(file_path)
+        
+        documents = loader.load()
+
+        # Split text
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        texts = text_splitter.split_documents(documents)
+
+        # Create or update vector store
+        if self.vector_store is None:
+            self.vector_store = FAISS.from_documents(texts, self.embeddings)
+        else:
+            self.vector_store.add_documents(texts)
+        
+        # Save index
+        self.vector_store.save_local(self.index_path)
+        print(f"Ingested {file_path} and updated index.")
+
+    def search(self, query: str, k: int = 3) -> List[str]:
+        if self.vector_store is None:
+            return []
+        
+        docs = self.vector_store.similarity_search(query, k=k)
+        return [doc.page_content for doc in docs]
+
+    def clear_index(self):
+        if os.path.exists(self.index_path):
+            import shutil
+            shutil.rmtree(self.index_path)
+        self.vector_store = None
+        print("Index cleared.")

requirements.txt

@@ -0,0 +1,23 @@
+torch
+torchvision
+torchaudio
+transformers==4.40.2
+diffusers==0.29.0
+accelerate
+protobuf
+sentencepiece
+fastapi
+uvicorn
+Pillow
+sqlalchemy
+passlib[bcrypt]
+python-jose[cryptography]
+python-multipart
+bcrypt
+duckduckgo-search
+langchain
+langchain-community
+sentence-transformers
+faiss-cpu
+pypdf
+python-multipart

schemas.py

@@ -0,0 +1,80 @@
+from pydantic import BaseModel
+from typing import Optional, List
+from datetime import datetime
+
+class UserBase(BaseModel):
+    email: str
+    full_name: Optional[str] = None
+    company_name: Optional[str] = None
+
+class UserCreate(UserBase):
+    password: str
+
+class User(UserBase):
+    id: int
+    is_admin: bool
+
+    class Config:
+        orm_mode = True
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+class TokenData(BaseModel):
+    email: Optional[str] = None
+
+class UserActivity(User):
+    message_count: int
+    prompt_count: int
+
+class ConversationBase(BaseModel):
+    title: str
+
+class ConversationCreate(ConversationBase):
+    pass
+
+class Conversation(ConversationBase):
+    id: int
+    user_id: int
+    created_at: datetime
+    updated_at: datetime
+
+    class Config:
+        orm_mode = True
+
+class ChatMessageBase(BaseModel):
+    role: str
+    content: str
+    conversation_id: Optional[int] = None
+
+class ChatMessageCreate(ChatMessageBase):
+    pass
+
+class ChatMessage(ChatMessageBase):
+    id: int
+    user_id: int
+    timestamp: datetime
+
+    class Config:
+        orm_mode = True
+
+class UserActivity(User):
+    message_count: int
+    prompt_count: int
+
+class SavedPromptBase(BaseModel):
+    title: str
+    content: str
+    is_public: bool = False
+
+class SavedPromptCreate(SavedPromptBase):
+    pass
+
+class SavedPrompt(SavedPromptBase):
+    id: int
+    user_id: int
+    created_at: datetime
+
+    class Config:
+        orm_mode = True

search_engine.py

@@ -0,0 +1,21 @@
+from duckduckgo_search import DDGS
+
+class SearchEngine:
+    def __init__(self):
+        self.ddgs = DDGS()
+
+    def search(self, query, max_results=3):
+        print(f"Searching web for: '{query}'")
+        try:
+            results = list(self.ddgs.text(query, max_results=max_results))
+            formatted_results = ""
+            for i, result in enumerate(results):
+                formatted_results += f"Source {i+1}: {result['title']}\nURL: {result['href']}\nContent: {result['body']}\n\n"
+            return formatted_results
+        except Exception as e:
+            print(f"Search failed: {e}")
+            return "Error: Could not perform search."
+
+if __name__ == "__main__":
+    se = SearchEngine()
+    print(se.search("What is the price of Bitcoin today?"))