mirror of
http://88.130.71.182:3000/BlitTech/contexta_be.git
synced 2026-06-12 23:23:21 +00:00
Initial commit
This commit is contained in:
belviskhoremk
0
app/__init__.py
Normal file
0
app/__init__.py
Normal file
127
app/config.py
Normal file
127
app/config.py
Normal file
@@ -0,0 +1,127 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
from typing import Optional, List
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
# App
|
||||
app_env: str = "development"
|
||||
app_secret_key: str = "dev-secret-key"
|
||||
allowed_origins: str = "http://localhost:5173,http://localhost:3000"
|
||||
|
||||
# Supabase
|
||||
supabase_url: str = ""
|
||||
supabase_anon_key: str = ""
|
||||
supabase_service_role_key: str = ""
|
||||
|
||||
# Qdrant
|
||||
qdrant_url: str = "http://localhost:6333"
|
||||
qdrant_api_key: Optional[str] = None
|
||||
|
||||
# LLM Providers
|
||||
openai_api_key: Optional[str] = None
|
||||
anthropic_api_key: Optional[str] = None
|
||||
google_api_key: Optional[str] = None
|
||||
fireworks_api_key: Optional[str] = None
|
||||
|
||||
# Embeddings
|
||||
embedding_model: str = "text-embedding-3-small"
|
||||
|
||||
# Stripe
|
||||
stripe_secret_key: str = ""
|
||||
stripe_webhook_secret: str = ""
|
||||
stripe_starter_price_id: str = ""
|
||||
stripe_pro_price_id: str = ""
|
||||
|
||||
# Redis
|
||||
redis_url: str = "redis://localhost:6379"
|
||||
|
||||
# Sentry
|
||||
sentry_dsn: Optional[str] = None
|
||||
|
||||
# Files
|
||||
max_file_size_mb: int = 50
|
||||
|
||||
@property
|
||||
def allowed_origins_list(self) -> List[str]:
|
||||
return [o.strip() for o in self.allowed_origins.split(",")]
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
settings = Settings()
|
||||
|
||||
# Plan limits
|
||||
PLAN_LIMITS = {
|
||||
"free": {
|
||||
"max_chatbots": 999999, # unlimited creation
|
||||
"max_published": 0, # cannot publish
|
||||
"models": [],
|
||||
"conversations_limit": 999999, # unlimited preview
|
||||
"code_export": False,
|
||||
"features": ["preview_mode", "testing"],
|
||||
},
|
||||
"starter": {
|
||||
"max_chatbots": 999999,
|
||||
"max_published": 1,
|
||||
"models": [
|
||||
"accounts/fireworks/models/llama-v3p1-70b-instruct",
|
||||
"accounts/fireworks/models/mixtral-8x7b-instruct",
|
||||
"accounts/fireworks/models/qwen2p5-72b-instruct",
|
||||
],
|
||||
"conversations_limit": 5000,
|
||||
"code_export": False,
|
||||
"features": ["marketplace", "analytics", "branding"],
|
||||
},
|
||||
"pro": {
|
||||
"max_chatbots": 3,
|
||||
"max_published": 3,
|
||||
"models": [
|
||||
"accounts/fireworks/models/llama-v3p1-70b-instruct",
|
||||
"accounts/fireworks/models/mixtral-8x7b-instruct",
|
||||
"gpt-4o",
|
||||
"gpt-4-turbo",
|
||||
"gpt-3.5-turbo",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"claude-3-opus-20240229",
|
||||
"gemini-1.5-pro",
|
||||
],
|
||||
"conversations_limit": 20000,
|
||||
"code_export": True,
|
||||
"features": [
|
||||
"marketplace",
|
||||
"code_export",
|
||||
"advanced_analytics",
|
||||
"priority_support",
|
||||
"custom_domain",
|
||||
"ab_testing",
|
||||
],
|
||||
},
|
||||
"enterprise": {
|
||||
"max_chatbots": 999999,
|
||||
"max_published": 999999,
|
||||
"models": ["*"],
|
||||
"conversations_limit": 999999,
|
||||
"code_export": True,
|
||||
"features": ["*"],
|
||||
},
|
||||
}
|
||||
|
||||
MODEL_PROVIDERS = {
|
||||
"accounts/fireworks/models/llama-v3p1-70b-instruct": "fireworks",
|
||||
"accounts/fireworks/models/mixtral-8x7b-instruct": "fireworks",
|
||||
"accounts/fireworks/models/qwen2p5-72b-instruct": "fireworks",
|
||||
"gpt-4o": "openai",
|
||||
"gpt-4-turbo": "openai",
|
||||
"gpt-3.5-turbo": "openai",
|
||||
"claude-3-5-sonnet-20241022": "anthropic",
|
||||
"claude-3-opus-20240229": "anthropic",
|
||||
"gemini-1.5-pro": "google",
|
||||
}
|
||||
|
||||
DEFAULT_MODELS = {
|
||||
"starter": "accounts/fireworks/models/llama-v3p1-70b-instruct",
|
||||
"pro": "gpt-4o",
|
||||
"enterprise": "gpt-4o",
|
||||
}
|
||||
82
app/dependencies.py
Normal file
82
app/dependencies.py
Normal file
@@ -0,0 +1,82 @@
|
||||
from fastapi import Depends, HTTPException, status, Header
|
||||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
|
||||
from typing import Optional
|
||||
from app.database import get_supabase
|
||||
from app.config import settings
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
security = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
|
||||
):
|
||||
"""Extract and verify the current user from Supabase JWT"""
|
||||
if not credentials:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Not authenticated",
|
||||
)
|
||||
|
||||
token = credentials.credentials
|
||||
supabase = get_supabase()
|
||||
|
||||
try:
|
||||
response = supabase.auth.get_user(token)
|
||||
if not response or not response.user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid or expired token",
|
||||
)
|
||||
return response.user
|
||||
except Exception as e:
|
||||
logger.error(f"Auth error: {e}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid or expired token",
|
||||
)
|
||||
|
||||
|
||||
async def get_optional_user(
|
||||
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
|
||||
):
|
||||
"""Optional auth - returns None if not authenticated"""
|
||||
if not credentials:
|
||||
return None
|
||||
try:
|
||||
return await get_current_user(credentials)
|
||||
except HTTPException:
|
||||
return None
|
||||
|
||||
|
||||
async def get_user_subscription(user=Depends(get_current_user)):
|
||||
"""Get user's subscription plan"""
|
||||
supabase = get_supabase()
|
||||
try:
|
||||
result = (
|
||||
supabase.table("subscriptions")
|
||||
.select("*")
|
||||
.eq("user_id", user.id)
|
||||
.eq("status", "active")
|
||||
.execute()
|
||||
)
|
||||
if result.data:
|
||||
return result.data[0]
|
||||
return {"plan": "free", "status": "active", "user_id": user.id}
|
||||
except Exception:
|
||||
return {"plan": "free", "status": "active", "user_id": user.id}
|
||||
|
||||
|
||||
async def require_plan(min_plan: str, user=Depends(get_current_user)):
|
||||
"""Require a minimum plan level"""
|
||||
plan_order = ["free", "starter", "pro", "enterprise"]
|
||||
subscription = await get_user_subscription(user)
|
||||
user_plan = subscription.get("plan", "free")
|
||||
|
||||
if plan_order.index(user_plan) < plan_order.index(min_plan):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
||||
detail=f"This feature requires {min_plan} plan or higher",
|
||||
)
|
||||
return user
|
||||
74
app/main.py
Normal file
74
app/main.py
Normal file
@@ -0,0 +1,74 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse
|
||||
import logging
|
||||
|
||||
from app.config import settings
|
||||
from app.routers import auth, chatbots, documents, chat, marketplace, billing
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── App ────────────────────────────────────────────────────────────────────────
|
||||
app = FastAPI(
|
||||
title="Contexta API",
|
||||
description="AI Chatbot Platform - Create, deploy and share custom AI chatbots powered by your data",
|
||||
version="1.0.0",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc",
|
||||
)
|
||||
|
||||
# ── Middleware ─────────────────────────────────────────────────────────────────
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.allowed_origins_list,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# ── Routers ────────────────────────────────────────────────────────────────────
|
||||
app.include_router(auth.router, prefix="/api/v1")
|
||||
app.include_router(chatbots.router, prefix="/api/v1")
|
||||
app.include_router(documents.router, prefix="/api/v1")
|
||||
app.include_router(chat.router, prefix="/api/v1")
|
||||
app.include_router(marketplace.router, prefix="/api/v1")
|
||||
app.include_router(billing.router, prefix="/api/v1")
|
||||
|
||||
# ── Health & Info ──────────────────────────────────────────────────────────────
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {
|
||||
"name": "Contexta API",
|
||||
"version": "1.0.0",
|
||||
"status": "running",
|
||||
"docs": "/docs",
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "healthy", "environment": settings.app_env}
|
||||
|
||||
|
||||
# ── Sentry ────────────────────────────────────────────────────────────────────
|
||||
if settings.sentry_dsn:
|
||||
import sentry_sdk
|
||||
sentry_sdk.init(dsn=settings.sentry_dsn, traces_sample_rate=0.1)
|
||||
logger.info("Sentry initialized")
|
||||
|
||||
# ── Startup ───────────────────────────────────────────────────────────────────
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
logger.info("Contexta API starting up...")
|
||||
logger.info(f"Environment: {settings.app_env}")
|
||||
logger.info(f"Allowed origins: {settings.allowed_origins_list}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)
|
||||
306
app/models.py
Normal file
306
app/models.py
Normal file
@@ -0,0 +1,306 @@
|
||||
from pydantic import BaseModel, EmailStr, Field
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
import uuid
|
||||
|
||||
|
||||
# ─── Enums ────────────────────────────────────────────────────────────────────
|
||||
|
||||
class PlanType(str, Enum):
|
||||
free = "free"
|
||||
starter = "starter"
|
||||
pro = "pro"
|
||||
enterprise = "enterprise"
|
||||
|
||||
|
||||
class SubscriptionStatus(str, Enum):
|
||||
active = "active"
|
||||
canceled = "canceled"
|
||||
past_due = "past_due"
|
||||
unpaid = "unpaid"
|
||||
trialing = "trialing"
|
||||
|
||||
|
||||
class ChatbotVisibility(str, Enum):
|
||||
preview = "preview"
|
||||
published = "published"
|
||||
|
||||
|
||||
class DocumentStatus(str, Enum):
|
||||
pending = "pending"
|
||||
processing = "processing"
|
||||
completed = "completed"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
user = "user"
|
||||
assistant = "assistant"
|
||||
system = "system"
|
||||
|
||||
|
||||
# ─── Auth Models ──────────────────────────────────────────────────────────────
|
||||
|
||||
class UserSignup(BaseModel):
|
||||
email: EmailStr
|
||||
password: str = Field(min_length=8)
|
||||
company_name: str = Field(min_length=2, max_length=100)
|
||||
|
||||
|
||||
class UserLogin(BaseModel):
|
||||
email: EmailStr
|
||||
password: str
|
||||
|
||||
|
||||
class UserResponse(BaseModel):
|
||||
id: str
|
||||
email: str
|
||||
company_name: Optional[str] = None
|
||||
plan: str = "free"
|
||||
created_at: Optional[datetime] = None
|
||||
|
||||
|
||||
class TokenResponse(BaseModel):
|
||||
access_token: str
|
||||
token_type: str = "bearer"
|
||||
user: UserResponse
|
||||
|
||||
|
||||
# ─── Company Models ────────────────────────────────────────────────────────────
|
||||
|
||||
class CompanyCreate(BaseModel):
|
||||
name: str = Field(min_length=2, max_length=100)
|
||||
website: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
|
||||
|
||||
class CompanyUpdate(BaseModel):
|
||||
name: Optional[str] = None
|
||||
website: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
logo_url: Optional[str] = None
|
||||
|
||||
|
||||
class CompanyResponse(BaseModel):
|
||||
id: str
|
||||
owner_id: str
|
||||
name: str
|
||||
website: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
logo_url: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
|
||||
|
||||
# ─── Chatbot Models ────────────────────────────────────────────────────────────
|
||||
|
||||
class ChatbotCreate(BaseModel):
|
||||
name: str = Field(min_length=2, max_length=100)
|
||||
description: Optional[str] = None
|
||||
system_prompt: Optional[str] = None
|
||||
model: str = "accounts/fireworks/models/llama-v3p1-70b-instruct"
|
||||
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
|
||||
max_tokens: int = Field(default=1000, ge=100, le=8000)
|
||||
primary_color: str = "#6366f1"
|
||||
welcome_message: str = "Hello! How can I help you today?"
|
||||
category: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
languages: List[str] = ["en"]
|
||||
|
||||
|
||||
class ChatbotUpdate(BaseModel):
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
system_prompt: Optional[str] = None
|
||||
model: Optional[str] = None
|
||||
temperature: Optional[float] = None
|
||||
max_tokens: Optional[int] = None
|
||||
primary_color: Optional[str] = None
|
||||
welcome_message: Optional[str] = None
|
||||
category: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
languages: Optional[List[str]] = None
|
||||
|
||||
|
||||
class ChatbotResponse(BaseModel):
|
||||
id: str
|
||||
company_id: str
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
system_prompt: Optional[str] = None
|
||||
model: str
|
||||
temperature: float
|
||||
max_tokens: int
|
||||
primary_color: str
|
||||
welcome_message: str
|
||||
category: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
languages: List[str]
|
||||
visibility: str
|
||||
is_published: bool
|
||||
qdrant_collection_name: Optional[str] = None
|
||||
document_count: int = 0
|
||||
conversation_count: int = 0
|
||||
average_rating: Optional[float] = None
|
||||
created_at: Optional[datetime] = None
|
||||
published_at: Optional[datetime] = None
|
||||
|
||||
|
||||
class ChatbotPublicResponse(BaseModel):
|
||||
"""For marketplace display"""
|
||||
id: str
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
category: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
languages: List[str]
|
||||
primary_color: str
|
||||
welcome_message: str
|
||||
average_rating: Optional[float] = None
|
||||
total_conversations: int = 0
|
||||
company_name: Optional[str] = None
|
||||
company_logo: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
published_at: Optional[datetime] = None
|
||||
|
||||
|
||||
# ─── Document Models ───────────────────────────────────────────────────────────
|
||||
|
||||
class DocumentResponse(BaseModel):
|
||||
id: str
|
||||
chatbot_id: str
|
||||
file_name: str
|
||||
file_type: str
|
||||
file_size: int
|
||||
chunk_count: int
|
||||
status: str
|
||||
error_message: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
|
||||
|
||||
# ─── Chat Models ───────────────────────────────────────────────────────────────
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
message: str = Field(min_length=1, max_length=4000)
|
||||
session_id: Optional[str] = None
|
||||
language: str = "en"
|
||||
|
||||
|
||||
class SourceDocument(BaseModel):
|
||||
document_name: str
|
||||
chunk_text: str
|
||||
score: float
|
||||
page_number: Optional[int] = None
|
||||
|
||||
|
||||
class ChatResponse(BaseModel):
|
||||
response: str
|
||||
session_id: str
|
||||
sources: List[SourceDocument] = []
|
||||
model_used: str
|
||||
tokens_used: int = 0
|
||||
|
||||
|
||||
class MessageResponse(BaseModel):
|
||||
id: str
|
||||
role: str
|
||||
content: str
|
||||
sources: Optional[List[Dict]] = None
|
||||
created_at: Optional[datetime] = None
|
||||
|
||||
|
||||
class ConversationResponse(BaseModel):
|
||||
id: str
|
||||
chatbot_id: str
|
||||
session_id: Optional[str] = None
|
||||
language: str
|
||||
message_count: int
|
||||
created_at: Optional[datetime] = None
|
||||
messages: List[MessageResponse] = []
|
||||
|
||||
|
||||
# ─── Subscription Models ───────────────────────────────────────────────────────
|
||||
|
||||
class SubscriptionResponse(BaseModel):
|
||||
id: str
|
||||
user_id: str
|
||||
plan: str
|
||||
status: str
|
||||
stripe_customer_id: Optional[str] = None
|
||||
current_period_start: Optional[datetime] = None
|
||||
current_period_end: Optional[datetime] = None
|
||||
chatbots_published: int = 0
|
||||
conversations_used: int = 0
|
||||
created_at: Optional[datetime] = None
|
||||
|
||||
|
||||
class CheckoutSessionCreate(BaseModel):
|
||||
plan: str # starter or pro
|
||||
success_url: str
|
||||
cancel_url: str
|
||||
|
||||
|
||||
class CheckoutSessionResponse(BaseModel):
|
||||
checkout_url: str
|
||||
session_id: str
|
||||
|
||||
|
||||
# ─── Analytics Models ──────────────────────────────────────────────────────────
|
||||
|
||||
class ChatbotAnalytics(BaseModel):
|
||||
chatbot_id: str
|
||||
total_conversations: int = 0
|
||||
unique_users: int = 0
|
||||
average_conversation_length: float = 0.0
|
||||
total_messages: int = 0
|
||||
average_rating: float = 0.0
|
||||
top_queries: List[str] = []
|
||||
conversations_last_7_days: List[Dict] = []
|
||||
conversations_last_30_days: int = 0
|
||||
|
||||
|
||||
# ─── Marketplace Models ────────────────────────────────────────────────────────
|
||||
|
||||
class MarketplaceFilter(BaseModel):
|
||||
category: Optional[str] = None
|
||||
industry: Optional[str] = None
|
||||
language: Optional[str] = None
|
||||
search: Optional[str] = None
|
||||
page: int = 1
|
||||
limit: int = 20
|
||||
|
||||
|
||||
class MarketplaceResponse(BaseModel):
|
||||
chatbots: List[ChatbotPublicResponse]
|
||||
total: int
|
||||
page: int
|
||||
limit: int
|
||||
has_more: bool
|
||||
|
||||
|
||||
class RatingCreate(BaseModel):
|
||||
rating: int = Field(ge=1, le=5)
|
||||
feedback: Optional[str] = None
|
||||
|
||||
|
||||
# ─── Code Export Models ────────────────────────────────────────────────────────
|
||||
|
||||
class CodeExportRequest(BaseModel):
|
||||
chatbot_id: str
|
||||
include_frontend: bool = True
|
||||
|
||||
|
||||
# ─── Generic Response Models ───────────────────────────────────────────────────
|
||||
|
||||
class MessageDetail(BaseModel):
|
||||
detail: str
|
||||
|
||||
|
||||
class SuccessResponse(BaseModel):
|
||||
success: bool
|
||||
message: str
|
||||
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
error: str
|
||||
detail: Optional[str] = None
|
||||
0
app/routers/__init__.py
Normal file
0
app/routers/__init__.py
Normal file
133
app/routers/auth.py
Normal file
133
app/routers/auth.py
Normal file
@@ -0,0 +1,133 @@
|
||||
from fastapi import APIRouter, HTTPException, status, Depends
|
||||
from app.models import UserSignup, UserLogin, UserResponse, TokenResponse
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_current_user
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/auth", tags=["Authentication"])
|
||||
|
||||
|
||||
@router.post("/signup", response_model=TokenResponse)
|
||||
async def signup(data: UserSignup):
|
||||
supabase = get_supabase()
|
||||
try:
|
||||
# Create auth user
|
||||
auth_resp = supabase.auth.sign_up(
|
||||
{"email": data.email, "password": data.password}
|
||||
)
|
||||
if not auth_resp.user:
|
||||
raise HTTPException(status_code=400, detail="Failed to create account")
|
||||
|
||||
user = auth_resp.user
|
||||
|
||||
# Create company record
|
||||
supabase.table("companies").insert(
|
||||
{
|
||||
"owner_id": user.id,
|
||||
"name": data.company_name,
|
||||
}
|
||||
).execute()
|
||||
|
||||
# Create free subscription
|
||||
supabase.table("subscriptions").insert(
|
||||
{
|
||||
"user_id": user.id,
|
||||
"plan": "free",
|
||||
"status": "active",
|
||||
}
|
||||
).execute()
|
||||
|
||||
token = auth_resp.session.access_token if auth_resp.session else ""
|
||||
return TokenResponse(
|
||||
access_token=token,
|
||||
user=UserResponse(
|
||||
id=user.id,
|
||||
email=user.email,
|
||||
company_name=data.company_name,
|
||||
plan="free",
|
||||
),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Signup error: {e}")
|
||||
if "already registered" in str(e).lower() or "already exists" in str(e).lower():
|
||||
raise HTTPException(status_code=400, detail="Email already registered")
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/login", response_model=TokenResponse)
|
||||
async def login(data: UserLogin):
|
||||
supabase = get_supabase()
|
||||
try:
|
||||
auth_resp = supabase.auth.sign_in_with_password(
|
||||
{"email": data.email, "password": data.password}
|
||||
)
|
||||
if not auth_resp.user or not auth_resp.session:
|
||||
raise HTTPException(status_code=401, detail="Invalid credentials")
|
||||
|
||||
user = auth_resp.user
|
||||
|
||||
# Get company info
|
||||
company = supabase.table("companies").select("name").eq("owner_id", user.id).execute()
|
||||
company_name = company.data[0]["name"] if company.data else ""
|
||||
|
||||
# Get subscription
|
||||
sub = (
|
||||
supabase.table("subscriptions")
|
||||
.select("plan")
|
||||
.eq("user_id", user.id)
|
||||
.eq("status", "active")
|
||||
.execute()
|
||||
)
|
||||
plan = sub.data[0]["plan"] if sub.data else "free"
|
||||
|
||||
return TokenResponse(
|
||||
access_token=auth_resp.session.access_token,
|
||||
user=UserResponse(
|
||||
id=user.id,
|
||||
email=user.email,
|
||||
company_name=company_name,
|
||||
plan=plan,
|
||||
),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Login error: {e}")
|
||||
raise HTTPException(status_code=401, detail="Invalid credentials")
|
||||
|
||||
|
||||
@router.post("/logout")
|
||||
async def logout(user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
try:
|
||||
supabase.auth.sign_out()
|
||||
except Exception:
|
||||
pass
|
||||
return {"message": "Logged out successfully"}
|
||||
|
||||
|
||||
@router.get("/me", response_model=UserResponse)
|
||||
async def get_me(user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
|
||||
company = supabase.table("companies").select("name").eq("owner_id", user.id).execute()
|
||||
company_name = company.data[0]["name"] if company.data else ""
|
||||
|
||||
sub = (
|
||||
supabase.table("subscriptions")
|
||||
.select("plan")
|
||||
.eq("user_id", user.id)
|
||||
.eq("status", "active")
|
||||
.execute()
|
||||
)
|
||||
plan = sub.data[0]["plan"] if sub.data else "free"
|
||||
|
||||
return UserResponse(
|
||||
id=user.id,
|
||||
email=user.email,
|
||||
company_name=company_name,
|
||||
plan=plan,
|
||||
)
|
||||
187
app/routers/billing.py
Normal file
187
app/routers/billing.py
Normal file
@@ -0,0 +1,187 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends, Request, Header
|
||||
from app.models import CheckoutSessionCreate, CheckoutSessionResponse, SubscriptionResponse
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_current_user
|
||||
from app.config import settings
|
||||
from typing import Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/billing", tags=["Billing"])
|
||||
|
||||
PLAN_PRICE_IDS = {
|
||||
"starter": settings.stripe_starter_price_id,
|
||||
"pro": settings.stripe_pro_price_id,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/checkout", response_model=CheckoutSessionResponse)
|
||||
async def create_checkout_session(data: CheckoutSessionCreate, user=Depends(get_current_user)):
|
||||
try:
|
||||
import stripe
|
||||
stripe.api_key = settings.stripe_secret_key
|
||||
|
||||
if data.plan not in PLAN_PRICE_IDS:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid plan: {data.plan}")
|
||||
|
||||
price_id = PLAN_PRICE_IDS[data.plan]
|
||||
if not price_id:
|
||||
raise HTTPException(status_code=400, detail="Plan price not configured")
|
||||
|
||||
supabase = get_supabase()
|
||||
sub = supabase.table("subscriptions").select("stripe_customer_id").eq("user_id", user.id).execute()
|
||||
|
||||
customer_id = None
|
||||
if sub.data and sub.data[0].get("stripe_customer_id"):
|
||||
customer_id = sub.data[0]["stripe_customer_id"]
|
||||
else:
|
||||
customer = stripe.Customer.create(email=user.email)
|
||||
customer_id = customer.id
|
||||
|
||||
session = stripe.checkout.Session.create(
|
||||
customer=customer_id,
|
||||
payment_method_types=["card"],
|
||||
line_items=[{"price": price_id, "quantity": 1}],
|
||||
mode="subscription",
|
||||
success_url=data.success_url + "?session_id={CHECKOUT_SESSION_ID}",
|
||||
cancel_url=data.cancel_url,
|
||||
metadata={"user_id": user.id, "plan": data.plan},
|
||||
)
|
||||
|
||||
return CheckoutSessionResponse(checkout_url=session.url, session_id=session.id)
|
||||
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="Stripe not configured")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Checkout error: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/webhook")
|
||||
async def stripe_webhook(
|
||||
request: Request,
|
||||
stripe_signature: Optional[str] = Header(None),
|
||||
):
|
||||
try:
|
||||
import stripe
|
||||
|
||||
stripe.api_key = settings.stripe_secret_key
|
||||
payload = await request.body()
|
||||
|
||||
if settings.stripe_webhook_secret and stripe_signature:
|
||||
try:
|
||||
event = stripe.Webhook.construct_event(
|
||||
payload, stripe_signature, settings.stripe_webhook_secret
|
||||
)
|
||||
except stripe.error.SignatureVerificationError:
|
||||
raise HTTPException(status_code=400, detail="Invalid signature")
|
||||
else:
|
||||
import json
|
||||
event = json.loads(payload)
|
||||
|
||||
supabase = get_supabase()
|
||||
event_type = event.get("type", "")
|
||||
|
||||
if event_type == "checkout.session.completed":
|
||||
session = event["data"]["object"]
|
||||
user_id = session.get("metadata", {}).get("user_id")
|
||||
plan = session.get("metadata", {}).get("plan", "starter")
|
||||
customer_id = session.get("customer")
|
||||
subscription_id = session.get("subscription")
|
||||
|
||||
if user_id:
|
||||
supabase.table("subscriptions").upsert({
|
||||
"user_id": user_id,
|
||||
"plan": plan,
|
||||
"status": "active",
|
||||
"stripe_customer_id": customer_id,
|
||||
"stripe_subscription_id": subscription_id,
|
||||
}, on_conflict="user_id").execute()
|
||||
|
||||
elif event_type in ("customer.subscription.updated", "customer.subscription.deleted"):
|
||||
sub_obj = event["data"]["object"]
|
||||
customer_id = sub_obj.get("customer")
|
||||
status = sub_obj.get("status", "canceled")
|
||||
|
||||
existing = supabase.table("subscriptions").select("*").eq("stripe_customer_id", customer_id).execute()
|
||||
if existing.data:
|
||||
mapped_status = "active" if status in ("active", "trialing") else "canceled"
|
||||
supabase.table("subscriptions").update({
|
||||
"status": mapped_status,
|
||||
}).eq("stripe_customer_id", customer_id).execute()
|
||||
|
||||
# Unpublish chatbots if subscription canceled
|
||||
if mapped_status == "canceled":
|
||||
user_id = existing.data[0]["user_id"]
|
||||
company = supabase.table("companies").select("id").eq("owner_id", user_id).execute()
|
||||
if company.data:
|
||||
supabase.table("chatbots").update({
|
||||
"is_published": False,
|
||||
"visibility": "preview",
|
||||
}).eq("company_id", company.data[0]["id"]).execute()
|
||||
|
||||
return {"received": True}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Webhook error: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/subscription", response_model=SubscriptionResponse)
|
||||
async def get_subscription(user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
result = supabase.table("subscriptions").select("*").eq("user_id", user.id).execute()
|
||||
|
||||
if not result.data:
|
||||
return SubscriptionResponse(
|
||||
id="free",
|
||||
user_id=user.id,
|
||||
plan="free",
|
||||
status="active",
|
||||
)
|
||||
|
||||
sub = result.data[0]
|
||||
return SubscriptionResponse(
|
||||
id=sub["id"],
|
||||
user_id=sub["user_id"],
|
||||
plan=sub["plan"],
|
||||
status=sub["status"],
|
||||
stripe_customer_id=sub.get("stripe_customer_id"),
|
||||
current_period_start=sub.get("current_period_start"),
|
||||
current_period_end=sub.get("current_period_end"),
|
||||
chatbots_published=sub.get("chatbots_published", 0),
|
||||
conversations_used=sub.get("conversations_used", 0),
|
||||
created_at=sub.get("created_at"),
|
||||
)
|
||||
|
||||
|
||||
@router.post("/portal")
|
||||
async def customer_portal(request: Request, user=Depends(get_current_user)):
|
||||
"""Create Stripe customer portal session"""
|
||||
try:
|
||||
import stripe
|
||||
stripe.api_key = settings.stripe_secret_key
|
||||
|
||||
supabase = get_supabase()
|
||||
sub = supabase.table("subscriptions").select("stripe_customer_id").eq("user_id", user.id).execute()
|
||||
|
||||
if not sub.data or not sub.data[0].get("stripe_customer_id"):
|
||||
raise HTTPException(status_code=404, detail="No subscription found")
|
||||
|
||||
body = await request.json()
|
||||
return_url = body.get("return_url", "http://localhost:5173/settings")
|
||||
|
||||
session = stripe.billing_portal.Session.create(
|
||||
customer=sub.data[0]["stripe_customer_id"],
|
||||
return_url=return_url,
|
||||
)
|
||||
return {"url": session.url}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
212
app/routers/chat.py
Normal file
212
app/routers/chat.py
Normal file
@@ -0,0 +1,212 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from app.models import ChatMessage, ChatResponse, ConversationResponse, MessageResponse
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_current_user, get_optional_user
|
||||
from app.services.rag import rag_engine
|
||||
from typing import List, Optional
|
||||
import uuid
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(tags=["Chat"])
|
||||
|
||||
|
||||
def _get_public_chatbot(chatbot_id: str, supabase) -> dict:
|
||||
"""Get a published chatbot (or any chatbot for preview)"""
|
||||
result = supabase.table("chatbots").select("*, companies(name, logo_url)").eq("id", chatbot_id).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found")
|
||||
return result.data[0]
|
||||
|
||||
|
||||
@router.post("/chat/{chatbot_id}", response_model=ChatResponse)
|
||||
async def chat(
|
||||
chatbot_id: str,
|
||||
message: ChatMessage,
|
||||
user=Depends(get_optional_user),
|
||||
):
|
||||
supabase = get_supabase()
|
||||
chatbot = _get_public_chatbot(chatbot_id, supabase)
|
||||
|
||||
# Allow preview access for owner, require published for public
|
||||
if not chatbot.get("is_published"):
|
||||
if not user:
|
||||
raise HTTPException(status_code=403, detail="This chatbot is in preview mode")
|
||||
# Check ownership
|
||||
company = supabase.table("companies").select("id").eq("owner_id", user.id).execute()
|
||||
if not company.data or company.data[0]["id"] != chatbot.get("company_id"):
|
||||
raise HTTPException(status_code=403, detail="This chatbot is in preview mode")
|
||||
|
||||
collection_name = chatbot.get("qdrant_collection_name")
|
||||
if not collection_name:
|
||||
raise HTTPException(status_code=400, detail="Chatbot has no knowledge base configured")
|
||||
|
||||
# Get or create conversation
|
||||
session_id = message.session_id or str(uuid.uuid4())
|
||||
conversation = _get_or_create_conversation(
|
||||
chatbot_id=chatbot_id,
|
||||
session_id=session_id,
|
||||
user_id=user.id if user else None,
|
||||
language=message.language,
|
||||
supabase=supabase,
|
||||
)
|
||||
|
||||
# Get conversation history
|
||||
history = _get_conversation_history(conversation["id"], supabase)
|
||||
|
||||
# Get company info for context
|
||||
company_data = chatbot.get("companies", {}) or {}
|
||||
chatbot_config = {
|
||||
**chatbot,
|
||||
"company_name": company_data.get("name", ""),
|
||||
}
|
||||
|
||||
# Run RAG
|
||||
result = await rag_engine.process_query(
|
||||
query=message.message,
|
||||
collection_name=collection_name,
|
||||
chatbot_config=chatbot_config,
|
||||
conversation_history=history,
|
||||
language=message.language,
|
||||
)
|
||||
|
||||
# Save messages
|
||||
_save_message(conversation["id"], "user", message.message, supabase)
|
||||
_save_message(
|
||||
conversation["id"],
|
||||
"assistant",
|
||||
result["response"],
|
||||
supabase,
|
||||
sources=[s.model_dump() for s in result.get("sources", [])],
|
||||
model=result.get("model", ""),
|
||||
)
|
||||
|
||||
# Update conversation message count
|
||||
supabase.table("conversations").update({
|
||||
"message_count": len(history) + 2
|
||||
}).eq("id", conversation["id"]).execute()
|
||||
|
||||
return ChatResponse(
|
||||
response=result["response"],
|
||||
session_id=session_id,
|
||||
sources=result.get("sources", []),
|
||||
model_used=result.get("model", ""),
|
||||
tokens_used=result.get("tokens_used", 0),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chat/{chatbot_id}/history/{session_id}", response_model=List[MessageResponse])
|
||||
async def get_chat_history(
|
||||
chatbot_id: str,
|
||||
session_id: str,
|
||||
user=Depends(get_optional_user),
|
||||
):
|
||||
supabase = get_supabase()
|
||||
|
||||
conversation = supabase.table("conversations").select("*") \
|
||||
.eq("chatbot_id", chatbot_id) \
|
||||
.eq("session_id", session_id) \
|
||||
.execute()
|
||||
|
||||
if not conversation.data:
|
||||
return []
|
||||
|
||||
conv_id = conversation.data[0]["id"]
|
||||
messages = supabase.table("messages").select("*") \
|
||||
.eq("conversation_id", conv_id) \
|
||||
.order("created_at", asc=True) \
|
||||
.execute()
|
||||
|
||||
return [
|
||||
MessageResponse(
|
||||
id=m["id"],
|
||||
role=m["role"],
|
||||
content=m["content"],
|
||||
sources=m.get("sources"),
|
||||
created_at=m.get("created_at"),
|
||||
)
|
||||
for m in (messages.data or [])
|
||||
]
|
||||
|
||||
|
||||
# ── Analytics endpoint ────────────────────────────────────────────────────────
|
||||
|
||||
@router.get("/analytics/{chatbot_id}")
|
||||
async def get_analytics(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
|
||||
# Verify ownership
|
||||
company = supabase.table("companies").select("id").eq("owner_id", user.id).execute()
|
||||
if not company.data:
|
||||
raise HTTPException(status_code=404, detail="Company not found")
|
||||
|
||||
chatbot = supabase.table("chatbots").select("id").eq("id", chatbot_id).eq("company_id", company.data[0]["id"]).execute()
|
||||
if not chatbot.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found")
|
||||
|
||||
total_convs = supabase.table("conversations").select("id", count="exact").eq("chatbot_id", chatbot_id).execute()
|
||||
total_msgs = supabase.table("messages").select("id", count="exact").execute()
|
||||
|
||||
return {
|
||||
"chatbot_id": chatbot_id,
|
||||
"total_conversations": total_convs.count or 0,
|
||||
"total_messages": total_msgs.count or 0,
|
||||
"average_rating": 0.0,
|
||||
"conversations_last_30_days": total_convs.count or 0,
|
||||
}
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _get_or_create_conversation(
|
||||
chatbot_id: str,
|
||||
session_id: str,
|
||||
user_id: Optional[str],
|
||||
language: str,
|
||||
supabase,
|
||||
) -> dict:
|
||||
existing = supabase.table("conversations").select("*") \
|
||||
.eq("chatbot_id", chatbot_id) \
|
||||
.eq("session_id", session_id) \
|
||||
.execute()
|
||||
|
||||
if existing.data:
|
||||
return existing.data[0]
|
||||
|
||||
new_conv = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"chatbot_id": chatbot_id,
|
||||
"user_id": user_id,
|
||||
"session_id": session_id,
|
||||
"language": language,
|
||||
"message_count": 0,
|
||||
}
|
||||
result = supabase.table("conversations").insert(new_conv).execute()
|
||||
return result.data[0]
|
||||
|
||||
|
||||
def _get_conversation_history(conversation_id: str, supabase) -> List[dict]:
|
||||
messages = supabase.table("messages").select("role, content") \
|
||||
.eq("conversation_id", conversation_id) \
|
||||
.order("created_at", asc=True) \
|
||||
.limit(20) \
|
||||
.execute()
|
||||
return messages.data or []
|
||||
|
||||
|
||||
def _save_message(
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
supabase,
|
||||
sources: Optional[list] = None,
|
||||
model: str = "",
|
||||
):
|
||||
supabase.table("messages").insert({
|
||||
"id": str(uuid.uuid4()),
|
||||
"conversation_id": conversation_id,
|
||||
"role": role,
|
||||
"content": content,
|
||||
"sources": sources,
|
||||
"model": model,
|
||||
}).execute()
|
||||
240
app/routers/chatbots.py
Normal file
240
app/routers/chatbots.py
Normal file
@@ -0,0 +1,240 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends, status
|
||||
from app.models import (
|
||||
ChatbotCreate, ChatbotUpdate, ChatbotResponse, SuccessResponse
|
||||
)
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_current_user, get_user_subscription
|
||||
from app.services.vector_store import vector_store
|
||||
from app.config import PLAN_LIMITS
|
||||
from typing import List
|
||||
import uuid
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/chatbots", tags=["Chatbots"])
|
||||
|
||||
|
||||
def _get_user_company(user_id: str, supabase) -> dict:
|
||||
result = supabase.table("companies").select("*").eq("owner_id", user_id).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=404, detail="Company not found")
|
||||
return result.data[0]
|
||||
|
||||
|
||||
async def _check_plan_limits(user_id: str, supabase, action: str = "create"):
|
||||
sub = supabase.table("subscriptions").select("*").eq("user_id", user_id).eq("status", "active").execute()
|
||||
plan = sub.data[0]["plan"] if sub.data else "free"
|
||||
limits = PLAN_LIMITS[plan]
|
||||
|
||||
if action == "publish":
|
||||
published = supabase.table("chatbots").select("id", count="exact") \
|
||||
.eq("company_id", _get_user_company(user_id, supabase)["id"]) \
|
||||
.eq("is_published", True).execute()
|
||||
count = published.count or 0
|
||||
max_pub = limits["max_published"]
|
||||
if max_pub == 0:
|
||||
raise HTTPException(status_code=402, detail="Upgrade to publish chatbots to marketplace")
|
||||
if count >= max_pub:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail=f"Publish limit reached ({max_pub}). Upgrade to publish more chatbots."
|
||||
)
|
||||
return plan
|
||||
|
||||
|
||||
@router.post("", response_model=ChatbotResponse, status_code=201)
|
||||
async def create_chatbot(data: ChatbotCreate, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
|
||||
# Create Qdrant collection
|
||||
collection_name = f"company_{company['id']}_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
vector_store.create_collection(collection_name)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create Qdrant collection: {e}")
|
||||
# Continue without vector store for now
|
||||
collection_name = None
|
||||
|
||||
chatbot_data = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"company_id": company["id"],
|
||||
"name": data.name,
|
||||
"description": data.description,
|
||||
"system_prompt": data.system_prompt,
|
||||
"model": data.model,
|
||||
"temperature": data.temperature,
|
||||
"max_tokens": data.max_tokens,
|
||||
"primary_color": data.primary_color,
|
||||
"welcome_message": data.welcome_message,
|
||||
"category": data.category,
|
||||
"industry": data.industry,
|
||||
"languages": data.languages,
|
||||
"visibility": "preview",
|
||||
"is_published": False,
|
||||
"qdrant_collection_name": collection_name,
|
||||
}
|
||||
|
||||
result = supabase.table("chatbots").insert(chatbot_data).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=500, detail="Failed to create chatbot")
|
||||
|
||||
return _format_chatbot(result.data[0], supabase)
|
||||
|
||||
|
||||
@router.get("", response_model=List[ChatbotResponse])
|
||||
async def list_chatbots(user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
|
||||
result = supabase.table("chatbots").select("*") \
|
||||
.eq("company_id", company["id"]) \
|
||||
.order("created_at", desc=True) \
|
||||
.execute()
|
||||
|
||||
return [_format_chatbot(c, supabase) for c in (result.data or [])]
|
||||
|
||||
|
||||
@router.get("/{chatbot_id}", response_model=ChatbotResponse)
|
||||
async def get_chatbot(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
chatbot = _get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
return _format_chatbot(chatbot, supabase)
|
||||
|
||||
|
||||
@router.put("/{chatbot_id}", response_model=ChatbotResponse)
|
||||
async def update_chatbot(chatbot_id: str, data: ChatbotUpdate, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
_get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
|
||||
update_data = {k: v for k, v in data.model_dump().items() if v is not None}
|
||||
if not update_data:
|
||||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
|
||||
result = supabase.table("chatbots").update(update_data).eq("id", chatbot_id).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=500, detail="Update failed")
|
||||
|
||||
return _format_chatbot(result.data[0], supabase)
|
||||
|
||||
|
||||
@router.delete("/{chatbot_id}", response_model=SuccessResponse)
|
||||
async def delete_chatbot(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
chatbot = _get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
|
||||
# Delete Qdrant collection
|
||||
if chatbot.get("qdrant_collection_name"):
|
||||
try:
|
||||
vector_store.delete_collection(chatbot["qdrant_collection_name"])
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to delete collection: {e}")
|
||||
|
||||
supabase.table("chatbots").delete().eq("id", chatbot_id).execute()
|
||||
return SuccessResponse(success=True, message="Chatbot deleted")
|
||||
|
||||
|
||||
@router.post("/{chatbot_id}/publish", response_model=ChatbotResponse)
|
||||
async def publish_chatbot(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
chatbot = _get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
|
||||
await _check_plan_limits(user.id, supabase, "publish")
|
||||
|
||||
result = supabase.table("chatbots").update({
|
||||
"is_published": True,
|
||||
"visibility": "published",
|
||||
}).eq("id", chatbot_id).execute()
|
||||
|
||||
return _format_chatbot(result.data[0], supabase)
|
||||
|
||||
|
||||
@router.post("/{chatbot_id}/unpublish", response_model=ChatbotResponse)
|
||||
async def unpublish_chatbot(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
_get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
|
||||
result = supabase.table("chatbots").update({
|
||||
"is_published": False,
|
||||
"visibility": "preview",
|
||||
}).eq("id", chatbot_id).execute()
|
||||
|
||||
return _format_chatbot(result.data[0], supabase)
|
||||
|
||||
|
||||
@router.post("/{chatbot_id}/export")
|
||||
async def export_chatbot(chatbot_id: str, user=Depends(get_current_user)):
|
||||
from fastapi.responses import StreamingResponse
|
||||
from app.services.code_export import generate_export_package
|
||||
from app.config import settings
|
||||
|
||||
supabase = get_supabase()
|
||||
company = _get_user_company(user.id, supabase)
|
||||
chatbot = _get_owned_chatbot(chatbot_id, company["id"], supabase)
|
||||
|
||||
# Check plan
|
||||
sub = supabase.table("subscriptions").select("plan").eq("user_id", user.id).eq("status", "active").execute()
|
||||
plan = sub.data[0]["plan"] if sub.data else "free"
|
||||
if plan not in ("pro", "enterprise"):
|
||||
raise HTTPException(status_code=402, detail="Code export requires Pro plan or higher")
|
||||
|
||||
zip_bytes = generate_export_package(
|
||||
chatbot=chatbot,
|
||||
company=company,
|
||||
qdrant_url=settings.qdrant_url,
|
||||
qdrant_key=settings.qdrant_api_key or "",
|
||||
)
|
||||
|
||||
filename = chatbot["name"].lower().replace(" ", "-") + "-chatbot.zip"
|
||||
return StreamingResponse(
|
||||
iter([zip_bytes]),
|
||||
media_type="application/zip",
|
||||
headers={"Content-Disposition": f"attachment; filename={filename}"},
|
||||
)
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _get_owned_chatbot(chatbot_id: str, company_id: str, supabase) -> dict:
|
||||
result = supabase.table("chatbots").select("*").eq("id", chatbot_id).eq("company_id", company_id).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found")
|
||||
return result.data[0]
|
||||
|
||||
|
||||
def _format_chatbot(chatbot: dict, supabase) -> ChatbotResponse:
|
||||
doc_count = supabase.table("documents").select("id", count="exact") \
|
||||
.eq("chatbot_id", chatbot["id"]) \
|
||||
.eq("status", "completed") \
|
||||
.execute()
|
||||
conv_count = supabase.table("conversations").select("id", count="exact") \
|
||||
.eq("chatbot_id", chatbot["id"]) \
|
||||
.execute()
|
||||
|
||||
return ChatbotResponse(
|
||||
id=chatbot["id"],
|
||||
company_id=chatbot["company_id"],
|
||||
name=chatbot["name"],
|
||||
description=chatbot.get("description"),
|
||||
system_prompt=chatbot.get("system_prompt"),
|
||||
model=chatbot.get("model", "accounts/fireworks/models/llama-v3p1-70b-instruct"),
|
||||
temperature=chatbot.get("temperature", 0.7),
|
||||
max_tokens=chatbot.get("max_tokens", 1000),
|
||||
primary_color=chatbot.get("primary_color", "#6366f1"),
|
||||
welcome_message=chatbot.get("welcome_message", "Hello! How can I help?"),
|
||||
category=chatbot.get("category"),
|
||||
industry=chatbot.get("industry"),
|
||||
languages=chatbot.get("languages", ["en"]),
|
||||
visibility=chatbot.get("visibility", "preview"),
|
||||
is_published=chatbot.get("is_published", False),
|
||||
qdrant_collection_name=chatbot.get("qdrant_collection_name"),
|
||||
document_count=doc_count.count or 0,
|
||||
conversation_count=conv_count.count or 0,
|
||||
created_at=chatbot.get("created_at"),
|
||||
published_at=chatbot.get("published_at"),
|
||||
)
|
||||
208
app/routers/documents.py
Normal file
208
app/routers/documents.py
Normal file
@@ -0,0 +1,208 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends, UploadFile, File, BackgroundTasks
|
||||
from app.models import DocumentResponse, SuccessResponse
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_current_user
|
||||
from app.services.document_processor import process_document
|
||||
from app.services.embeddings import embedding_service
|
||||
from app.services.vector_store import vector_store
|
||||
from app.config import settings
|
||||
from typing import List
|
||||
import uuid
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/chatbots/{chatbot_id}/documents", tags=["Documents"])
|
||||
|
||||
ALLOWED_TYPES = {
|
||||
"application/pdf": ".pdf",
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
|
||||
"text/csv": ".csv",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
|
||||
"text/plain": ".txt",
|
||||
"text/markdown": ".md",
|
||||
}
|
||||
|
||||
ALLOWED_EXTENSIONS = {".pdf", ".docx", ".csv", ".xlsx", ".txt", ".md"}
|
||||
|
||||
|
||||
def _get_user_chatbot(chatbot_id: str, user_id: str, supabase) -> dict:
|
||||
company = supabase.table("companies").select("id").eq("owner_id", user_id).execute()
|
||||
if not company.data:
|
||||
raise HTTPException(status_code=404, detail="Company not found")
|
||||
company_id = company.data[0]["id"]
|
||||
|
||||
chatbot = supabase.table("chatbots").select("*").eq("id", chatbot_id).eq("company_id", company_id).execute()
|
||||
if not chatbot.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found")
|
||||
return chatbot.data[0]
|
||||
|
||||
|
||||
@router.post("", response_model=DocumentResponse, status_code=201)
|
||||
async def upload_document(
|
||||
chatbot_id: str,
|
||||
background_tasks: BackgroundTasks,
|
||||
file: UploadFile = File(...),
|
||||
user=Depends(get_current_user),
|
||||
):
|
||||
supabase = get_supabase()
|
||||
chatbot = _get_user_chatbot(chatbot_id, user.id, supabase)
|
||||
|
||||
# Validate file
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="No filename provided")
|
||||
|
||||
ext = "." + file.filename.rsplit(".", 1)[-1].lower() if "." in file.filename else ""
|
||||
if ext not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"File type not supported. Allowed: {', '.join(ALLOWED_EXTENSIONS)}"
|
||||
)
|
||||
|
||||
# Read file
|
||||
file_bytes = await file.read()
|
||||
file_size = len(file_bytes)
|
||||
max_size = settings.max_file_size_mb * 1024 * 1024
|
||||
|
||||
if file_size > max_size:
|
||||
raise HTTPException(
|
||||
status_code=413,
|
||||
detail=f"File too large. Max size: {settings.max_file_size_mb}MB"
|
||||
)
|
||||
|
||||
# Create document record
|
||||
doc_id = str(uuid.uuid4())
|
||||
doc_data = {
|
||||
"id": doc_id,
|
||||
"chatbot_id": chatbot_id,
|
||||
"file_name": file.filename,
|
||||
"file_type": ext,
|
||||
"file_size": file_size,
|
||||
"chunk_count": 0,
|
||||
"status": "processing",
|
||||
}
|
||||
|
||||
result = supabase.table("documents").insert(doc_data).execute()
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=500, detail="Failed to create document record")
|
||||
|
||||
# Process in background
|
||||
background_tasks.add_task(
|
||||
_process_document_bg,
|
||||
file_bytes=file_bytes,
|
||||
file_name=file.filename,
|
||||
doc_id=doc_id,
|
||||
chatbot=chatbot,
|
||||
supabase=supabase,
|
||||
)
|
||||
|
||||
return DocumentResponse(**result.data[0])
|
||||
|
||||
|
||||
async def _process_document_bg(
|
||||
file_bytes: bytes,
|
||||
file_name: str,
|
||||
doc_id: str,
|
||||
chatbot: dict,
|
||||
supabase,
|
||||
):
|
||||
"""Background task to process and embed a document"""
|
||||
try:
|
||||
company_id = chatbot.get("company_id", "")
|
||||
collection_name = chatbot.get("qdrant_collection_name")
|
||||
|
||||
if not collection_name:
|
||||
logger.error(f"No Qdrant collection for chatbot {chatbot['id']}")
|
||||
supabase.table("documents").update({
|
||||
"status": "failed",
|
||||
"error_message": "Vector store not configured"
|
||||
}).eq("id", doc_id).execute()
|
||||
return
|
||||
|
||||
# Ensure collection exists
|
||||
if not vector_store.collection_exists(collection_name):
|
||||
vector_store.create_collection(collection_name)
|
||||
|
||||
# Process document
|
||||
chunks, payloads = process_document(
|
||||
file_bytes=file_bytes,
|
||||
file_name=file_name,
|
||||
document_id=doc_id,
|
||||
company_id=company_id,
|
||||
)
|
||||
|
||||
if not chunks:
|
||||
raise ValueError("No text extracted from document")
|
||||
|
||||
# Generate embeddings in batches
|
||||
batch_size = 50
|
||||
all_ids = []
|
||||
all_vectors = []
|
||||
all_payloads = []
|
||||
|
||||
for i in range(0, len(chunks), batch_size):
|
||||
batch_chunks = chunks[i:i + batch_size]
|
||||
batch_payloads = payloads[i:i + batch_size]
|
||||
|
||||
vectors = embedding_service.embed_batch(batch_chunks)
|
||||
ids = [str(uuid.uuid4()) for _ in vectors]
|
||||
|
||||
all_ids.extend(ids)
|
||||
all_vectors.extend(vectors)
|
||||
all_payloads.extend(batch_payloads)
|
||||
|
||||
# Upsert to Qdrant
|
||||
vector_store.upsert_vectors(
|
||||
collection_name=collection_name,
|
||||
vectors=all_vectors,
|
||||
payloads=all_payloads,
|
||||
ids=all_ids,
|
||||
)
|
||||
|
||||
# Update document record
|
||||
supabase.table("documents").update({
|
||||
"status": "completed",
|
||||
"chunk_count": len(chunks),
|
||||
}).eq("id", doc_id).execute()
|
||||
|
||||
logger.info(f"Document {doc_id} processed: {len(chunks)} chunks")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Document processing error for {doc_id}: {e}")
|
||||
supabase.table("documents").update({
|
||||
"status": "failed",
|
||||
"error_message": str(e)[:500],
|
||||
}).eq("id", doc_id).execute()
|
||||
|
||||
|
||||
@router.get("", response_model=List[DocumentResponse])
|
||||
async def list_documents(chatbot_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
_get_user_chatbot(chatbot_id, user.id, supabase)
|
||||
|
||||
result = supabase.table("documents").select("*") \
|
||||
.eq("chatbot_id", chatbot_id) \
|
||||
.order("created_at", desc=True) \
|
||||
.execute()
|
||||
|
||||
return [DocumentResponse(**d) for d in (result.data or [])]
|
||||
|
||||
|
||||
@router.delete("/{document_id}", response_model=SuccessResponse)
|
||||
async def delete_document(chatbot_id: str, document_id: str, user=Depends(get_current_user)):
|
||||
supabase = get_supabase()
|
||||
chatbot = _get_user_chatbot(chatbot_id, user.id, supabase)
|
||||
|
||||
doc = supabase.table("documents").select("*").eq("id", document_id).eq("chatbot_id", chatbot_id).execute()
|
||||
if not doc.data:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
|
||||
# Remove vectors from Qdrant
|
||||
collection_name = chatbot.get("qdrant_collection_name")
|
||||
if collection_name:
|
||||
try:
|
||||
vector_store.delete_by_document_id(collection_name, document_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to delete vectors: {e}")
|
||||
|
||||
supabase.table("documents").delete().eq("id", document_id).execute()
|
||||
return SuccessResponse(success=True, message="Document deleted")
|
||||
133
app/routers/marketplace.py
Normal file
133
app/routers/marketplace.py
Normal file
@@ -0,0 +1,133 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends, Query
|
||||
from app.models import ChatbotPublicResponse, MarketplaceResponse, RatingCreate
|
||||
from app.database import get_supabase
|
||||
from app.dependencies import get_optional_user, get_current_user
|
||||
from typing import Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/marketplace", tags=["Marketplace"])
|
||||
|
||||
CATEGORIES = [
|
||||
"Customer Support", "Sales", "FAQ", "E-commerce",
|
||||
"Healthcare", "Finance", "Education", "HR", "Legal", "Other"
|
||||
]
|
||||
|
||||
INDUSTRIES = [
|
||||
"Technology", "E-commerce", "Healthcare", "Finance",
|
||||
"Education", "Legal", "Real Estate", "Hospitality", "Retail", "Other"
|
||||
]
|
||||
|
||||
|
||||
@router.get("/chatbots", response_model=MarketplaceResponse)
|
||||
async def list_marketplace_chatbots(
|
||||
category: Optional[str] = Query(None),
|
||||
industry: Optional[str] = Query(None),
|
||||
language: Optional[str] = Query(None),
|
||||
search: Optional[str] = Query(None),
|
||||
page: int = Query(1, ge=1),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
user=Depends(get_optional_user),
|
||||
):
|
||||
supabase = get_supabase()
|
||||
|
||||
query = supabase.table("chatbots").select(
|
||||
"*, companies(name, logo_url)"
|
||||
).eq("is_published", True).eq("visibility", "published")
|
||||
|
||||
if category:
|
||||
query = query.eq("category", category)
|
||||
if industry:
|
||||
query = query.eq("industry", industry)
|
||||
if search:
|
||||
query = query.ilike("name", f"%{search}%")
|
||||
|
||||
offset = (page - 1) * limit
|
||||
result = query.order("created_at", desc=True).range(offset, offset + limit - 1).execute()
|
||||
all_result = supabase.table("chatbots").select("id", count="exact").eq("is_published", True).execute()
|
||||
total = all_result.count or 0
|
||||
|
||||
chatbots = []
|
||||
for c in (result.data or []):
|
||||
company_data = c.get("companies") or {}
|
||||
chatbots.append(
|
||||
ChatbotPublicResponse(
|
||||
id=c["id"],
|
||||
name=c["name"],
|
||||
description=c.get("description"),
|
||||
category=c.get("category"),
|
||||
industry=c.get("industry"),
|
||||
languages=c.get("languages", ["en"]),
|
||||
primary_color=c.get("primary_color", "#6366f1"),
|
||||
welcome_message=c.get("welcome_message", "Hello!"),
|
||||
average_rating=c.get("average_rating"),
|
||||
total_conversations=c.get("total_conversations", 0),
|
||||
company_name=company_data.get("name"),
|
||||
company_logo=company_data.get("logo_url"),
|
||||
created_at=c.get("created_at"),
|
||||
published_at=c.get("published_at"),
|
||||
)
|
||||
)
|
||||
|
||||
return MarketplaceResponse(
|
||||
chatbots=chatbots,
|
||||
total=total,
|
||||
page=page,
|
||||
limit=limit,
|
||||
has_more=(offset + limit) < total,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chatbots/{chatbot_id}", response_model=ChatbotPublicResponse)
|
||||
async def get_marketplace_chatbot(chatbot_id: str):
|
||||
supabase = get_supabase()
|
||||
result = supabase.table("chatbots").select("*, companies(name, logo_url)") \
|
||||
.eq("id", chatbot_id) \
|
||||
.eq("is_published", True) \
|
||||
.execute()
|
||||
|
||||
if not result.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found in marketplace")
|
||||
|
||||
c = result.data[0]
|
||||
company_data = c.get("companies") or {}
|
||||
return ChatbotPublicResponse(
|
||||
id=c["id"],
|
||||
name=c["name"],
|
||||
description=c.get("description"),
|
||||
category=c.get("category"),
|
||||
industry=c.get("industry"),
|
||||
languages=c.get("languages", ["en"]),
|
||||
primary_color=c.get("primary_color", "#6366f1"),
|
||||
welcome_message=c.get("welcome_message", "Hello!"),
|
||||
average_rating=c.get("average_rating"),
|
||||
total_conversations=c.get("total_conversations", 0),
|
||||
company_name=company_data.get("name"),
|
||||
company_logo=company_data.get("logo_url"),
|
||||
created_at=c.get("created_at"),
|
||||
published_at=c.get("published_at"),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/categories")
|
||||
async def get_categories():
|
||||
return {"categories": CATEGORIES, "industries": INDUSTRIES}
|
||||
|
||||
|
||||
@router.post("/chatbots/{chatbot_id}/rate")
|
||||
async def rate_chatbot(
|
||||
chatbot_id: str,
|
||||
rating: RatingCreate,
|
||||
user=Depends(get_current_user),
|
||||
):
|
||||
supabase = get_supabase()
|
||||
chatbot = supabase.table("chatbots").select("id, average_rating").eq("id", chatbot_id).eq("is_published", True).execute()
|
||||
if not chatbot.data:
|
||||
raise HTTPException(status_code=404, detail="Chatbot not found")
|
||||
|
||||
# Simple rating update (average)
|
||||
current = chatbot.data[0].get("average_rating") or rating.rating
|
||||
new_avg = (current + rating.rating) / 2
|
||||
|
||||
supabase.table("chatbots").update({"average_rating": round(new_avg, 1)}).eq("id", chatbot_id).execute()
|
||||
return {"message": "Rating submitted", "new_average": round(new_avg, 1)}
|
||||
0
app/services/__init__.py
Normal file
0
app/services/__init__.py
Normal file
713
app/services/code_export.py
Normal file
713
app/services/code_export.py
Normal file
@@ -0,0 +1,713 @@
|
||||
import zipfile
|
||||
import io
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
def generate_export_package(
|
||||
chatbot: Dict[str, Any],
|
||||
company: Dict[str, Any],
|
||||
qdrant_url: str,
|
||||
qdrant_key: str,
|
||||
) -> bytes:
|
||||
"""
|
||||
Generate a complete export ZIP with FastAPI backend + React widget
|
||||
"""
|
||||
buffer = io.BytesIO()
|
||||
|
||||
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
# ── Backend files ──────────────────────────────────────────
|
||||
zf.writestr("backend/requirements.txt", _requirements())
|
||||
zf.writestr("backend/.env.example", _env_example(chatbot, qdrant_url, qdrant_key))
|
||||
zf.writestr("backend/main.py", _main_py(chatbot))
|
||||
zf.writestr("backend/rag_engine.py", _rag_engine_py())
|
||||
zf.writestr("backend/Dockerfile", _dockerfile())
|
||||
zf.writestr("backend/docker-compose.yml", _docker_compose(chatbot))
|
||||
zf.writestr("backend/README.md", _backend_readme(chatbot))
|
||||
|
||||
# ── Frontend files ─────────────────────────────────────────
|
||||
zf.writestr("frontend/src/ChatWidget.tsx", _chat_widget_tsx(chatbot))
|
||||
zf.writestr("frontend/src/useChat.ts", _use_chat_ts())
|
||||
zf.writestr("frontend/src/api.ts", _api_ts())
|
||||
zf.writestr("frontend/src/types.ts", _types_ts())
|
||||
zf.writestr("frontend/package.json", _package_json(chatbot))
|
||||
zf.writestr("frontend/tsconfig.json", _tsconfig())
|
||||
zf.writestr("frontend/vite.config.ts", _vite_config())
|
||||
zf.writestr("frontend/README.md", _frontend_readme(chatbot))
|
||||
|
||||
# ── Root ───────────────────────────────────────────────────
|
||||
zf.writestr("QUICK_START.md", _quick_start(chatbot))
|
||||
zf.writestr("setup.py", _setup_wizard(chatbot))
|
||||
|
||||
buffer.seek(0)
|
||||
return buffer.read()
|
||||
|
||||
|
||||
def _requirements():
|
||||
return """fastapi==0.115.0
|
||||
uvicorn[standard]==0.30.6
|
||||
python-dotenv==1.0.1
|
||||
pydantic==2.8.2
|
||||
qdrant-client==1.11.1
|
||||
openai==1.51.0
|
||||
anthropic==0.34.2
|
||||
google-generativeai==0.8.1
|
||||
httpx==0.27.2
|
||||
langdetect==1.0.9
|
||||
"""
|
||||
|
||||
|
||||
def _env_example(chatbot: Dict, qdrant_url: str, qdrant_key: str):
|
||||
name = chatbot.get("name", "My Chatbot").upper().replace(" ", "_")
|
||||
return f"""# {name} - Environment Configuration
|
||||
# Copy to .env and fill in your values
|
||||
|
||||
# LLM Provider (choose one)
|
||||
LLM_PROVIDER=openai
|
||||
LLM_MODEL=gpt-4o
|
||||
LLM_API_KEY=sk-your-openai-key
|
||||
|
||||
# For Anthropic: sk-ant-your-key
|
||||
# For Google: your-google-api-key
|
||||
# For Fireworks: your-fireworks-key
|
||||
|
||||
# Embeddings (required - OpenAI)
|
||||
EMBEDDING_API_KEY=sk-your-openai-key
|
||||
EMBEDDING_MODEL=text-embedding-3-small
|
||||
|
||||
# Qdrant Vector Database
|
||||
QDRANT_URL={qdrant_url}
|
||||
QDRANT_API_KEY={qdrant_key}
|
||||
QDRANT_COLLECTION={chatbot.get("qdrant_collection_name", "chatbot_collection")}
|
||||
|
||||
# Server
|
||||
PORT=8000
|
||||
HOST=0.0.0.0
|
||||
ALLOWED_ORIGINS=http://localhost:3000,https://yourdomain.com
|
||||
"""
|
||||
|
||||
|
||||
def _main_py(chatbot: Dict):
|
||||
return f'''"""
|
||||
Auto-generated FastAPI backend for: {chatbot.get("name", "Chatbot")}
|
||||
Generated by Contexta Platform
|
||||
"""
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Optional
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from rag_engine import RAGEngine
|
||||
|
||||
load_dotenv()
|
||||
|
||||
app = FastAPI(
|
||||
title="{chatbot.get("name", "Chatbot")} API",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=os.getenv("ALLOWED_ORIGINS", "*").split(","),
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
rag = RAGEngine(
|
||||
qdrant_url=os.getenv("QDRANT_URL"),
|
||||
qdrant_api_key=os.getenv("QDRANT_API_KEY"),
|
||||
collection_name=os.getenv("QDRANT_COLLECTION"),
|
||||
llm_provider=os.getenv("LLM_PROVIDER", "openai"),
|
||||
llm_model=os.getenv("LLM_MODEL", "gpt-4o"),
|
||||
llm_api_key=os.getenv("LLM_API_KEY"),
|
||||
embedding_api_key=os.getenv("EMBEDDING_API_KEY"),
|
||||
embedding_model=os.getenv("EMBEDDING_MODEL", "text-embedding-3-small"),
|
||||
system_prompt="""{chatbot.get("system_prompt") or "You are a helpful assistant."}""",
|
||||
)
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
message: str
|
||||
session_id: Optional[str] = None
|
||||
language: str = "en"
|
||||
history: List[dict] = []
|
||||
|
||||
|
||||
class Source(BaseModel):
|
||||
document_name: str
|
||||
text: str
|
||||
score: float
|
||||
|
||||
|
||||
class ChatResponse(BaseModel):
|
||||
response: str
|
||||
session_id: str
|
||||
sources: List[Source]
|
||||
|
||||
|
||||
@app.post("/chat", response_model=ChatResponse)
|
||||
async def chat(request: ChatRequest):
|
||||
import uuid
|
||||
session_id = request.session_id or str(uuid.uuid4())
|
||||
result = await rag.query(
|
||||
query=request.message,
|
||||
history=request.history,
|
||||
language=request.language,
|
||||
)
|
||||
return ChatResponse(
|
||||
response=result["response"],
|
||||
session_id=session_id,
|
||||
sources=[Source(**s) for s in result.get("sources", [])],
|
||||
)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {{"status": "healthy", "chatbot": "{chatbot.get("name", "Chatbot")}"}}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host=os.getenv("HOST", "0.0.0.0"), port=int(os.getenv("PORT", 8000)))
|
||||
'''
|
||||
|
||||
|
||||
def _rag_engine_py():
|
||||
return '''"""RAG Engine - Retrieval-Augmented Generation"""
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.models import Distance, VectorParams
|
||||
from openai import AsyncOpenAI
|
||||
from typing import List, Dict, Any, Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RAGEngine:
|
||||
def __init__(self, qdrant_url, qdrant_api_key, collection_name,
|
||||
llm_provider, llm_model, llm_api_key,
|
||||
embedding_api_key, embedding_model, system_prompt=""):
|
||||
self.collection_name = collection_name
|
||||
self.llm_provider = llm_provider
|
||||
self.llm_model = llm_model
|
||||
self.llm_api_key = llm_api_key
|
||||
self.embedding_model = embedding_model
|
||||
self.system_prompt = system_prompt
|
||||
|
||||
# Qdrant
|
||||
qdrant_kwargs = {"url": qdrant_url}
|
||||
if qdrant_api_key:
|
||||
qdrant_kwargs["api_key"] = qdrant_api_key
|
||||
self.qdrant = QdrantClient(**qdrant_kwargs)
|
||||
|
||||
# OpenAI for embeddings
|
||||
self.embed_client = AsyncOpenAI(api_key=embedding_api_key)
|
||||
|
||||
async def embed(self, text: str) -> List[float]:
|
||||
resp = await self.embed_client.embeddings.create(
|
||||
model=self.embedding_model, input=text
|
||||
)
|
||||
return resp.data[0].embedding
|
||||
|
||||
async def retrieve(self, query_vector: List[float], limit: int = 5) -> List[Dict]:
|
||||
results = self.qdrant.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_vector,
|
||||
limit=limit,
|
||||
score_threshold=0.3,
|
||||
)
|
||||
return [{"text": r.payload.get("text", ""), "document_name": r.payload.get("file_name", ""), "score": r.score}
|
||||
for r in results]
|
||||
|
||||
async def generate(self, messages: List[Dict]) -> str:
|
||||
if self.llm_provider == "openai":
|
||||
from openai import AsyncOpenAI
|
||||
client = AsyncOpenAI(api_key=self.llm_api_key)
|
||||
resp = await client.chat.completions.create(
|
||||
model=self.llm_model, messages=messages, max_tokens=1000
|
||||
)
|
||||
return resp.choices[0].message.content
|
||||
elif self.llm_provider == "anthropic":
|
||||
import anthropic
|
||||
client = anthropic.AsyncAnthropic(api_key=self.llm_api_key)
|
||||
system = next((m["content"] for m in messages if m["role"] == "system"), "")
|
||||
conv = [m for m in messages if m["role"] != "system"]
|
||||
resp = await client.messages.create(
|
||||
model=self.llm_model, max_tokens=1000, system=system, messages=conv
|
||||
)
|
||||
return resp.content[0].text
|
||||
elif self.llm_provider == "fireworks":
|
||||
import httpx
|
||||
async with httpx.AsyncClient(timeout=60) as c:
|
||||
r = await c.post(
|
||||
"https://api.fireworks.ai/inference/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {self.llm_api_key}"},
|
||||
json={"model": self.llm_model, "messages": messages, "max_tokens": 1000},
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()["choices"][0]["message"]["content"]
|
||||
return "Error: unknown provider"
|
||||
|
||||
async def query(self, query: str, history: List[Dict] = None, language: str = "en") -> Dict:
|
||||
if history is None:
|
||||
history = []
|
||||
|
||||
query_vec = await self.embed(query)
|
||||
docs = await self.retrieve(query_vec)
|
||||
context = "\\n\\n---\\n\\n".join(d["text"] for d in docs) or "No context found."
|
||||
|
||||
system = f"{self.system_prompt}\\n\\nContext:\\n{context}"
|
||||
messages = [{"role": "system", "content": system}]
|
||||
for h in history[-10:]:
|
||||
messages.append(h)
|
||||
messages.append({"role": "user", "content": query})
|
||||
|
||||
response = await self.generate(messages)
|
||||
return {"response": response, "sources": docs}
|
||||
'''
|
||||
|
||||
|
||||
def _dockerfile():
|
||||
return """FROM python:3.11-slim
|
||||
WORKDIR /app
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
COPY . .
|
||||
EXPOSE 8000
|
||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
"""
|
||||
|
||||
|
||||
def _docker_compose(chatbot: Dict):
|
||||
name = chatbot.get("name", "chatbot").lower().replace(" ", "-")
|
||||
return f"""version: '3.8'
|
||||
services:
|
||||
api:
|
||||
build: .
|
||||
ports:
|
||||
- "8000:8000"
|
||||
env_file: .env
|
||||
restart: unless-stopped
|
||||
container_name: {name}-api
|
||||
"""
|
||||
|
||||
|
||||
def _chat_widget_tsx(chatbot: Dict):
|
||||
color = chatbot.get("primary_color", "#6366f1")
|
||||
welcome = chatbot.get("welcome_message", "Hello! How can I help you today?")
|
||||
name = chatbot.get("name", "Assistant")
|
||||
return f'''import React, {{ useState, useRef, useEffect }} from "react";
|
||||
import {{ useChat }} from "./useChat";
|
||||
|
||||
const PRIMARY_COLOR = "{color}";
|
||||
const BOT_NAME = "{name}";
|
||||
const WELCOME_MESSAGE = "{welcome}";
|
||||
|
||||
export const ChatWidget: React.FC = () => {{
|
||||
const [isOpen, setIsOpen] = useState(false);
|
||||
const {{ messages, isLoading, sendMessage }} = useChat(WELCOME_MESSAGE);
|
||||
const bottomRef = useRef<HTMLDivElement>(null);
|
||||
|
||||
useEffect(() => {{
|
||||
bottomRef.current?.scrollIntoView({{ behavior: "smooth" }});
|
||||
}}, [messages]);
|
||||
|
||||
return (
|
||||
<>
|
||||
{{isOpen && (
|
||||
<div style={{{{
|
||||
position: "fixed", bottom: 90, right: 20, width: 360, height: 520,
|
||||
borderRadius: 16, boxShadow: "0 20px 60px rgba(0,0,0,0.2)",
|
||||
display: "flex", flexDirection: "column", background: "#fff",
|
||||
fontFamily: "system-ui, -apple-system, sans-serif", zIndex: 9999
|
||||
}}}}>
|
||||
<div style={{{{ background: PRIMARY_COLOR, padding: "16px 20px",
|
||||
borderRadius: "16px 16px 0 0", display: "flex", justifyContent: "space-between", alignItems: "center" }}}}>
|
||||
<span style={{{{ color: "#fff", fontWeight: 600, fontSize: 16 }}}}>{{BOT_NAME}}</span>
|
||||
<button onClick={{() => setIsOpen(false)}}
|
||||
style={{{{ background: "none", border: "none", color: "#fff", cursor: "pointer", fontSize: 20 }}}}>×</button>
|
||||
</div>
|
||||
<div style={{{{ flex: 1, overflowY: "auto", padding: 16, display: "flex", flexDirection: "column", gap: 12 }}}}>
|
||||
{{messages.map((msg, i) => (
|
||||
<div key={{i}} style={{{{ display: "flex", justifyContent: msg.role === "user" ? "flex-end" : "flex-start" }}}}>
|
||||
<div style={{{{
|
||||
maxWidth: "80%", padding: "10px 14px", borderRadius: 12, fontSize: 14, lineHeight: 1.5,
|
||||
background: msg.role === "user" ? PRIMARY_COLOR : "#f3f4f6",
|
||||
color: msg.role === "user" ? "#fff" : "#111"
|
||||
}}}}>{{msg.content}}</div>
|
||||
</div>
|
||||
))}}
|
||||
{{isLoading && <div style={{{{ color: "#6b7280", fontSize: 13 }}}}>Thinking...</div>}}
|
||||
<div ref={{bottomRef}} />
|
||||
</div>
|
||||
<div style={{{{ padding: "12px 16px", borderTop: "1px solid #e5e7eb", display: "flex", gap: 8 }}}}>
|
||||
<input
|
||||
style={{{{ flex: 1, border: "1px solid #e5e7eb", borderRadius: 8, padding: "8px 12px", outline: "none", fontSize: 14 }}}}
|
||||
placeholder="Type a message..."
|
||||
onKeyDown={{(e) => {{
|
||||
if (e.key === "Enter" && !e.shiftKey) {{
|
||||
e.preventDefault();
|
||||
const val = (e.target as HTMLInputElement).value.trim();
|
||||
if (val) {{ sendMessage(val); (e.target as HTMLInputElement).value = ""; }}
|
||||
}}
|
||||
}}}}
|
||||
/>
|
||||
<button
|
||||
style={{{{ background: PRIMARY_COLOR, color: "#fff", border: "none", borderRadius: 8,
|
||||
padding: "8px 14px", cursor: "pointer", fontSize: 14 }}}}
|
||||
onClick={{(e) => {{
|
||||
const input = (e.currentTarget.previousSibling as HTMLInputElement);
|
||||
const val = input.value.trim();
|
||||
if (val) {{ sendMessage(val); input.value = ""; }}
|
||||
}}}}
|
||||
>Send</button>
|
||||
</div>
|
||||
</div>
|
||||
)}}
|
||||
<button
|
||||
onClick={{() => setIsOpen(!isOpen)}}
|
||||
style={{{{
|
||||
position: "fixed", bottom: 20, right: 20, width: 56, height: 56,
|
||||
borderRadius: "50%", background: PRIMARY_COLOR, border: "none",
|
||||
cursor: "pointer", display: "flex", alignItems: "center", justifyContent: "center",
|
||||
boxShadow: "0 4px 20px rgba(0,0,0,0.2)", zIndex: 9999, fontSize: 24
|
||||
}}}}
|
||||
>
|
||||
{{isOpen ? "×" : "💬"}}
|
||||
</button>
|
||||
</>
|
||||
);
|
||||
}};
|
||||
'''
|
||||
|
||||
|
||||
def _use_chat_ts():
|
||||
return '''import { useState, useCallback } from "react";
|
||||
import { sendChatMessage } from "./api";
|
||||
|
||||
interface Message {
|
||||
role: "user" | "assistant";
|
||||
content: string;
|
||||
}
|
||||
|
||||
export function useChat(welcomeMessage: string) {
|
||||
const [messages, setMessages] = useState<Message[]>([
|
||||
{ role: "assistant", content: welcomeMessage }
|
||||
]);
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
const [sessionId] = useState(() => crypto.randomUUID());
|
||||
|
||||
const sendMessage = useCallback(async (content: string) => {
|
||||
setMessages(prev => [...prev, { role: "user", content }]);
|
||||
setIsLoading(true);
|
||||
try {
|
||||
const history = messages.map(m => ({ role: m.role, content: m.content }));
|
||||
const result = await sendChatMessage({ message: content, session_id: sessionId, history });
|
||||
setMessages(prev => [...prev, { role: "assistant", content: result.response }]);
|
||||
} catch {
|
||||
setMessages(prev => [...prev, { role: "assistant", content: "Sorry, I encountered an error. Please try again." }]);
|
||||
} finally {
|
||||
setIsLoading(false);
|
||||
}
|
||||
}, [messages, sessionId]);
|
||||
|
||||
return { messages, isLoading, sendMessage };
|
||||
}
|
||||
'''
|
||||
|
||||
|
||||
def _api_ts():
|
||||
return '''const API_URL = import.meta.env.VITE_API_URL || "http://localhost:8000";
|
||||
|
||||
export async function sendChatMessage(payload: {
|
||||
message: string;
|
||||
session_id: string;
|
||||
history?: any[];
|
||||
}) {
|
||||
const response = await fetch(`${API_URL}/chat`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
if (!response.ok) throw new Error("Chat request failed");
|
||||
return response.json();
|
||||
}
|
||||
'''
|
||||
|
||||
|
||||
def _types_ts():
|
||||
return '''export interface Message {
|
||||
role: "user" | "assistant";
|
||||
content: string;
|
||||
}
|
||||
|
||||
export interface Source {
|
||||
document_name: string;
|
||||
text: string;
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface ChatResponse {
|
||||
response: string;
|
||||
session_id: string;
|
||||
sources: Source[];
|
||||
}
|
||||
'''
|
||||
|
||||
|
||||
def _package_json(chatbot: Dict):
|
||||
name = chatbot.get("name", "chatbot").lower().replace(" ", "-")
|
||||
return f'''{{
|
||||
"name": "{name}-widget",
|
||||
"version": "1.0.0",
|
||||
"scripts": {{
|
||||
"dev": "vite",
|
||||
"build": "tsc && vite build",
|
||||
"preview": "vite preview"
|
||||
}},
|
||||
"dependencies": {{
|
||||
"react": "^18.2.0",
|
||||
"react-dom": "^18.2.0"
|
||||
}},
|
||||
"devDependencies": {{
|
||||
"@types/react": "^18.2.0",
|
||||
"@types/react-dom": "^18.2.0",
|
||||
"typescript": "^5.0.0",
|
||||
"vite": "^5.0.0",
|
||||
"@vitejs/plugin-react": "^4.0.0"
|
||||
}}
|
||||
}}
|
||||
'''
|
||||
|
||||
|
||||
def _tsconfig():
|
||||
return '''{
|
||||
"compilerOptions": {
|
||||
"target": "ES2020",
|
||||
"lib": ["ES2020", "DOM"],
|
||||
"module": "ESNext",
|
||||
"moduleResolution": "bundler",
|
||||
"jsx": "react-jsx",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
|
||||
def _vite_config():
|
||||
return '''import { defineConfig } from "vite";
|
||||
import react from "@vitejs/plugin-react";
|
||||
|
||||
export default defineConfig({
|
||||
plugins: [react()],
|
||||
build: {
|
||||
lib: {
|
||||
entry: "src/main.tsx",
|
||||
name: "ChatWidget",
|
||||
fileName: "chatbot-widget"
|
||||
},
|
||||
rollupOptions: {
|
||||
external: ["react", "react-dom"],
|
||||
}
|
||||
}
|
||||
});
|
||||
'''
|
||||
|
||||
|
||||
def _backend_readme(chatbot: Dict):
|
||||
return f"""# {chatbot.get("name", "Chatbot")} - Backend API
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# Edit .env with your API keys
|
||||
pip install -r requirements.txt
|
||||
uvicorn main:app --reload --port 8000
|
||||
```
|
||||
|
||||
## Deploy with Docker
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# Edit .env
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
- `POST /chat` - Send a message
|
||||
- `GET /health` - Health check
|
||||
|
||||
## Environment Variables
|
||||
|
||||
See `.env.example` for all required variables.
|
||||
"""
|
||||
|
||||
|
||||
def _frontend_readme(chatbot: Dict):
|
||||
return f"""# {chatbot.get("name", "Chatbot")} - Chat Widget
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# Set VITE_API_URL to your backend URL
|
||||
npm install
|
||||
npm run dev
|
||||
```
|
||||
|
||||
## Build for Production
|
||||
|
||||
```bash
|
||||
npm run build
|
||||
```
|
||||
|
||||
## Embed in Any Website
|
||||
|
||||
```html
|
||||
<script src="path/to/dist/chatbot-widget.umd.cjs"></script>
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
- `VITE_API_URL` - Backend API URL (default: http://localhost:8000)
|
||||
"""
|
||||
|
||||
|
||||
def _quick_start(chatbot: Dict):
|
||||
return f"""# Quick Start - {chatbot.get("name", "Chatbot")}
|
||||
|
||||
Get your chatbot running in 5 minutes!
|
||||
|
||||
## Prerequisites
|
||||
- Python 3.11+
|
||||
- Node.js 18+
|
||||
- API key from OpenAI, Anthropic, or Google
|
||||
|
||||
## 1. Configure Environment (2 min)
|
||||
|
||||
Run the setup wizard:
|
||||
```bash
|
||||
python setup.py
|
||||
```
|
||||
|
||||
Or manually:
|
||||
```bash
|
||||
cd backend
|
||||
cp .env.example .env
|
||||
# Edit .env with your keys
|
||||
```
|
||||
|
||||
## 2. Start Backend (1 min)
|
||||
|
||||
```bash
|
||||
cd backend
|
||||
pip install -r requirements.txt
|
||||
uvicorn main:app --reload
|
||||
```
|
||||
|
||||
Backend runs at: http://localhost:8000
|
||||
|
||||
## 3. Start Frontend Widget (1 min)
|
||||
|
||||
```bash
|
||||
cd frontend
|
||||
npm install
|
||||
npm run dev
|
||||
```
|
||||
|
||||
Widget available at: http://localhost:3000
|
||||
|
||||
## 4. Embed in Your Website
|
||||
|
||||
After building (`npm run build`):
|
||||
```html
|
||||
<script src="dist/chatbot-widget.umd.cjs"></script>
|
||||
```
|
||||
|
||||
## Deploy
|
||||
|
||||
### Railway (Recommended)
|
||||
```bash
|
||||
railway init
|
||||
railway up
|
||||
```
|
||||
|
||||
### Docker
|
||||
```bash
|
||||
cd backend && docker-compose up -d
|
||||
```
|
||||
"""
|
||||
|
||||
|
||||
def _setup_wizard(chatbot: Dict):
|
||||
return f'''#!/usr/bin/env python3
|
||||
"""
|
||||
Interactive setup wizard for {chatbot.get("name", "Chatbot")}
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main():
|
||||
print("""
|
||||
╔══════════════════════════════════════╗
|
||||
║ {chatbot.get("name", "Chatbot")} Setup Wizard ║
|
||||
╚══════════════════════════════════════╝
|
||||
""")
|
||||
|
||||
print("Choose your LLM provider:")
|
||||
print("1. OpenAI (GPT-4o)")
|
||||
print("2. Anthropic (Claude)")
|
||||
print("3. Google (Gemini)")
|
||||
print("4. Fireworks AI (Free, open-source models)")
|
||||
choice = input("\\nEnter choice (1-4): ").strip()
|
||||
|
||||
providers = {{"1": "openai", "2": "anthropic", "3": "google", "4": "fireworks"}}
|
||||
models = {{"1": "gpt-4o", "2": "claude-3-5-sonnet-20241022", "3": "gemini-1.5-pro",
|
||||
"4": "accounts/fireworks/models/llama-v3p1-70b-instruct"}}
|
||||
|
||||
provider = providers.get(choice, "openai")
|
||||
model = models.get(choice, "gpt-4o")
|
||||
|
||||
api_key = input(f"Enter your {{provider}} API key: ").strip()
|
||||
|
||||
env_content = f"""LLM_PROVIDER={{provider}}
|
||||
LLM_MODEL={{model}}
|
||||
LLM_API_KEY={{api_key}}
|
||||
EMBEDDING_API_KEY={{api_key if provider == "openai" else input("Enter OpenAI key for embeddings: ").strip()}}
|
||||
EMBEDDING_MODEL=text-embedding-3-small
|
||||
QDRANT_URL={os.getenv("QDRANT_URL", "your-qdrant-url")}
|
||||
QDRANT_API_KEY={os.getenv("QDRANT_API_KEY", "your-qdrant-key")}
|
||||
QDRANT_COLLECTION={chatbot.get("qdrant_collection_name", "chatbot_collection")}
|
||||
"""
|
||||
|
||||
env_file = Path("backend/.env")
|
||||
env_file.write_text(env_content)
|
||||
print("\\n✅ .env file created!")
|
||||
|
||||
frontend_url = input("\\nBackend URL for frontend (default: http://localhost:8000): ").strip()
|
||||
if not frontend_url:
|
||||
frontend_url = "http://localhost:8000"
|
||||
|
||||
Path("frontend/.env").write_text(f"VITE_API_URL={{frontend_url}}\\n")
|
||||
print("✅ Frontend .env created!")
|
||||
|
||||
print("""
|
||||
\\n╔══════════════════════════════════════╗
|
||||
║ Setup Complete! 🎉 ║
|
||||
╠══════════════════════════════════════╣
|
||||
║ Backend: cd backend && uvicorn ║
|
||||
║ main:app --reload ║
|
||||
║ Frontend: cd frontend && npm dev ║
|
||||
╚══════════════════════════════════════╝
|
||||
""")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
'''
|
||||
221
app/services/document_processor.py
Normal file
221
app/services/document_processor.py
Normal file
@@ -0,0 +1,221 @@
|
||||
import io
|
||||
import logging
|
||||
from typing import List, Dict, Any, Tuple
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CHUNK_SIZE = 512 # tokens approximate (chars ÷ 4)
|
||||
CHUNK_OVERLAP = 50
|
||||
|
||||
|
||||
def parse_pdf(file_bytes: bytes) -> List[Dict[str, Any]]:
|
||||
"""Parse PDF and return list of {text, page_number}"""
|
||||
try:
|
||||
import pypdf
|
||||
|
||||
reader = pypdf.PdfReader(io.BytesIO(file_bytes))
|
||||
pages = []
|
||||
for i, page in enumerate(reader.pages):
|
||||
text = page.extract_text() or ""
|
||||
text = text.strip()
|
||||
if text:
|
||||
pages.append({"text": text, "page_number": i + 1})
|
||||
return pages
|
||||
except Exception as e:
|
||||
logger.error(f"PDF parse error: {e}")
|
||||
raise ValueError(f"Failed to parse PDF: {str(e)}")
|
||||
|
||||
|
||||
def parse_docx(file_bytes: bytes) -> List[Dict[str, Any]]:
|
||||
"""Parse DOCX and return list of {text, page_number}"""
|
||||
try:
|
||||
from docx import Document
|
||||
|
||||
doc = Document(io.BytesIO(file_bytes))
|
||||
sections = []
|
||||
current_text = []
|
||||
section_idx = 1
|
||||
|
||||
for para in doc.paragraphs:
|
||||
text = para.text.strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# New section on headings
|
||||
if para.style.name.startswith("Heading"):
|
||||
if current_text:
|
||||
sections.append(
|
||||
{"text": "\n".join(current_text), "page_number": section_idx}
|
||||
)
|
||||
current_text = []
|
||||
section_idx += 1
|
||||
current_text.append(text)
|
||||
|
||||
if current_text:
|
||||
sections.append({"text": "\n".join(current_text), "page_number": section_idx})
|
||||
|
||||
return sections if sections else [{"text": "", "page_number": 1}]
|
||||
except Exception as e:
|
||||
logger.error(f"DOCX parse error: {e}")
|
||||
raise ValueError(f"Failed to parse DOCX: {str(e)}")
|
||||
|
||||
|
||||
def parse_csv(file_bytes: bytes) -> List[Dict[str, Any]]:
|
||||
"""Parse CSV - each row becomes a chunk"""
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
df = pd.read_csv(io.BytesIO(file_bytes))
|
||||
columns = list(df.columns)
|
||||
chunks = []
|
||||
|
||||
# Process in batches of rows
|
||||
batch_size = 10
|
||||
for start in range(0, len(df), batch_size):
|
||||
batch = df.iloc[start : start + batch_size]
|
||||
rows_text = []
|
||||
for _, row in batch.iterrows():
|
||||
row_parts = [f"{col}: {val}" for col, val in zip(columns, row) if str(val) != "nan"]
|
||||
rows_text.append(" | ".join(row_parts))
|
||||
text = "\n".join(rows_text)
|
||||
chunks.append({"text": text, "page_number": (start // batch_size) + 1})
|
||||
|
||||
return chunks
|
||||
except Exception as e:
|
||||
logger.error(f"CSV parse error: {e}")
|
||||
raise ValueError(f"Failed to parse CSV: {str(e)}")
|
||||
|
||||
|
||||
def parse_xlsx(file_bytes: bytes) -> List[Dict[str, Any]]:
|
||||
"""Parse XLSX - each sheet becomes sections"""
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
xl = pd.ExcelFile(io.BytesIO(file_bytes))
|
||||
chunks = []
|
||||
page_num = 1
|
||||
|
||||
for sheet_name in xl.sheet_names:
|
||||
df = xl.parse(sheet_name)
|
||||
columns = list(df.columns)
|
||||
|
||||
batch_size = 10
|
||||
for start in range(0, len(df), batch_size):
|
||||
batch = df.iloc[start : start + batch_size]
|
||||
rows_text = [f"Sheet: {sheet_name}"]
|
||||
for _, row in batch.iterrows():
|
||||
row_parts = [
|
||||
f"{col}: {val}"
|
||||
for col, val in zip(columns, row)
|
||||
if str(val) not in ("nan", "NaT", "None")
|
||||
]
|
||||
if row_parts:
|
||||
rows_text.append(" | ".join(row_parts))
|
||||
text = "\n".join(rows_text)
|
||||
if text.strip():
|
||||
chunks.append({"text": text, "page_number": page_num})
|
||||
page_num += 1
|
||||
|
||||
return chunks
|
||||
except Exception as e:
|
||||
logger.error(f"XLSX parse error: {e}")
|
||||
raise ValueError(f"Failed to parse XLSX: {str(e)}")
|
||||
|
||||
|
||||
def parse_txt(file_bytes: bytes) -> List[Dict[str, Any]]:
|
||||
"""Parse plain text"""
|
||||
try:
|
||||
text = file_bytes.decode("utf-8", errors="ignore")
|
||||
# Split into sections by double newlines
|
||||
sections = [s.strip() for s in text.split("\n\n") if s.strip()]
|
||||
if not sections:
|
||||
sections = [text.strip()]
|
||||
return [{"text": s, "page_number": i + 1} for i, s in enumerate(sections)]
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to parse TXT: {str(e)}")
|
||||
|
||||
|
||||
def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
|
||||
"""Split text into overlapping chunks"""
|
||||
# Approximate token count: 1 token ≈ 4 chars
|
||||
char_size = chunk_size * 4
|
||||
char_overlap = overlap * 4
|
||||
|
||||
if len(text) <= char_size:
|
||||
return [text]
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
while start < len(text):
|
||||
end = min(start + char_size, len(text))
|
||||
|
||||
# Try to break at sentence boundary
|
||||
if end < len(text):
|
||||
for sep in [". ", "! ", "? ", "\n", " "]:
|
||||
pos = text.rfind(sep, start, end)
|
||||
if pos > start + char_size // 2:
|
||||
end = pos + len(sep)
|
||||
break
|
||||
|
||||
chunk = text[start:end].strip()
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
|
||||
start = end - char_overlap if end - char_overlap > start else end
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def process_document(
|
||||
file_bytes: bytes,
|
||||
file_name: str,
|
||||
document_id: str,
|
||||
company_id: str,
|
||||
) -> Tuple[List[str], List[Dict[str, Any]]]:
|
||||
"""
|
||||
Main entry point: parse and chunk a document.
|
||||
Returns (chunks_text, chunk_payloads)
|
||||
"""
|
||||
ext = Path(file_name).suffix.lower()
|
||||
|
||||
# Parse
|
||||
if ext == ".pdf":
|
||||
pages = parse_pdf(file_bytes)
|
||||
elif ext == ".docx":
|
||||
pages = parse_docx(file_bytes)
|
||||
elif ext == ".csv":
|
||||
pages = parse_csv(file_bytes)
|
||||
elif ext in (".xlsx", ".xls"):
|
||||
pages = parse_xlsx(file_bytes)
|
||||
elif ext in (".txt", ".md"):
|
||||
pages = parse_txt(file_bytes)
|
||||
else:
|
||||
raise ValueError(f"Unsupported file type: {ext}")
|
||||
|
||||
# Chunk
|
||||
all_chunks = []
|
||||
all_payloads = []
|
||||
|
||||
for page in pages:
|
||||
text = page["text"]
|
||||
page_num = page.get("page_number", 1)
|
||||
|
||||
chunks = chunk_text(text)
|
||||
for idx, chunk in enumerate(chunks):
|
||||
all_chunks.append(chunk)
|
||||
all_payloads.append(
|
||||
{
|
||||
"document_id": document_id,
|
||||
"company_id": company_id,
|
||||
"file_name": file_name,
|
||||
"page_number": page_num,
|
||||
"chunk_index": idx,
|
||||
"text": chunk,
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Processed {file_name}: {len(pages)} pages → {len(all_chunks)} chunks"
|
||||
)
|
||||
return all_chunks, all_payloads
|
||||
54
app/services/embeddings.py
Normal file
54
app/services/embeddings.py
Normal file
@@ -0,0 +1,54 @@
|
||||
from openai import OpenAI
|
||||
from app.config import settings
|
||||
from typing import List
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_openai_client = None
|
||||
|
||||
|
||||
def get_openai_client() -> OpenAI:
|
||||
global _openai_client
|
||||
if _openai_client is None:
|
||||
_openai_client = OpenAI(api_key=settings.openai_api_key)
|
||||
return _openai_client
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
def __init__(self):
|
||||
self.model = settings.embedding_model
|
||||
|
||||
def embed_text(self, text: str) -> List[float]:
|
||||
"""Generate embedding for a single text"""
|
||||
client = get_openai_client()
|
||||
try:
|
||||
response = client.embeddings.create(
|
||||
model=self.model,
|
||||
input=text,
|
||||
)
|
||||
return response.data[0].embedding
|
||||
except Exception as e:
|
||||
logger.error(f"Embedding error: {e}")
|
||||
raise
|
||||
|
||||
def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""Generate embeddings for multiple texts"""
|
||||
client = get_openai_client()
|
||||
try:
|
||||
# Clean texts
|
||||
cleaned = [t.replace("\n", " ").strip() for t in texts if t.strip()]
|
||||
if not cleaned:
|
||||
return []
|
||||
|
||||
response = client.embeddings.create(
|
||||
model=self.model,
|
||||
input=cleaned,
|
||||
)
|
||||
return [item.embedding for item in response.data]
|
||||
except Exception as e:
|
||||
logger.error(f"Batch embedding error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
embedding_service = EmbeddingService()
|
||||
171
app/services/llm.py
Normal file
171
app/services/llm.py
Normal file
@@ -0,0 +1,171 @@
|
||||
from app.config import settings, MODEL_PROVIDERS, PLAN_LIMITS
|
||||
from typing import List, Dict, Any, Optional, AsyncGenerator
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMService:
|
||||
"""Routes requests to appropriate LLM provider"""
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
model: str,
|
||||
max_tokens: int = 1000,
|
||||
temperature: float = 0.7,
|
||||
) -> Dict[str, Any]:
|
||||
"""Generate a response from the LLM"""
|
||||
provider = MODEL_PROVIDERS.get(model, "openai")
|
||||
|
||||
try:
|
||||
if provider == "fireworks":
|
||||
return await self._call_fireworks(messages, model, max_tokens, temperature)
|
||||
elif provider == "openai":
|
||||
return await self._call_openai(messages, model, max_tokens, temperature)
|
||||
elif provider == "anthropic":
|
||||
return await self._call_anthropic(messages, model, max_tokens, temperature)
|
||||
elif provider == "google":
|
||||
return await self._call_google(messages, model, max_tokens, temperature)
|
||||
else:
|
||||
return await self._call_openai(messages, model, max_tokens, temperature)
|
||||
except Exception as e:
|
||||
logger.error(f"LLM error ({provider}/{model}): {e}")
|
||||
# Fallback to a basic model if available
|
||||
if model != "accounts/fireworks/models/llama-v3p1-70b-instruct" and settings.fireworks_api_key:
|
||||
logger.info("Falling back to Fireworks AI")
|
||||
return await self._call_fireworks(
|
||||
messages,
|
||||
"accounts/fireworks/models/llama-v3p1-70b-instruct",
|
||||
max_tokens,
|
||||
temperature,
|
||||
)
|
||||
raise
|
||||
|
||||
async def _call_fireworks(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
) -> Dict[str, Any]:
|
||||
import httpx
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {settings.fireworks_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
resp = await client.post(
|
||||
"https://api.fireworks.ai/inference/v1/chat/completions",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
return {
|
||||
"content": data["choices"][0]["message"]["content"],
|
||||
"tokens_used": data.get("usage", {}).get("total_tokens", 0),
|
||||
"model": model,
|
||||
}
|
||||
|
||||
async def _call_openai(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
) -> Dict[str, Any]:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
client = AsyncOpenAI(api_key=settings.openai_api_key)
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
return {
|
||||
"content": response.choices[0].message.content,
|
||||
"tokens_used": response.usage.total_tokens if response.usage else 0,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
async def _call_anthropic(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
) -> Dict[str, Any]:
|
||||
import anthropic
|
||||
|
||||
client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)
|
||||
|
||||
# Separate system message from conversation
|
||||
system_msg = ""
|
||||
conv_messages = []
|
||||
for msg in messages:
|
||||
if msg["role"] == "system":
|
||||
system_msg = msg["content"]
|
||||
else:
|
||||
conv_messages.append(msg)
|
||||
|
||||
response = await client.messages.create(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
system=system_msg if system_msg else "You are a helpful assistant.",
|
||||
messages=conv_messages,
|
||||
temperature=temperature,
|
||||
)
|
||||
return {
|
||||
"content": response.content[0].text,
|
||||
"tokens_used": response.usage.input_tokens + response.usage.output_tokens,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
async def _call_google(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
) -> Dict[str, Any]:
|
||||
import google.generativeai as genai
|
||||
|
||||
genai.configure(api_key=settings.google_api_key)
|
||||
gemini_model = genai.GenerativeModel(model)
|
||||
|
||||
# Convert messages
|
||||
parts = []
|
||||
for msg in messages:
|
||||
role = "user" if msg["role"] in ("user", "system") else "model"
|
||||
parts.append({"role": role, "parts": [msg["content"]]})
|
||||
|
||||
# Use last message as prompt if only one
|
||||
if len(parts) == 1:
|
||||
response = await gemini_model.generate_content_async(
|
||||
parts[0]["parts"][0],
|
||||
generation_config={"max_output_tokens": max_tokens, "temperature": temperature},
|
||||
)
|
||||
else:
|
||||
chat = gemini_model.start_chat(history=parts[:-1])
|
||||
response = await chat.send_message_async(
|
||||
parts[-1]["parts"][0],
|
||||
generation_config={"max_output_tokens": max_tokens, "temperature": temperature},
|
||||
)
|
||||
|
||||
return {
|
||||
"content": response.text,
|
||||
"tokens_used": 0,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
|
||||
llm_service = LLMService()
|
||||
130
app/services/rag.py
Normal file
130
app/services/rag.py
Normal file
@@ -0,0 +1,130 @@
|
||||
from app.services.embeddings import embedding_service
|
||||
from app.services.vector_store import vector_store
|
||||
from app.services.llm import llm_service
|
||||
from app.models import SourceDocument
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RAG_SYSTEM_PROMPT = """You are a helpful AI assistant for {company_name}.
|
||||
Your role is to answer questions based on the provided context from company documents.
|
||||
|
||||
IMPORTANT RULES:
|
||||
1. Only answer based on the provided context
|
||||
2. If information is not in the context, say "I don't have information about that in my knowledge base"
|
||||
3. Be concise and helpful
|
||||
4. Always maintain a professional, friendly tone
|
||||
5. If asked about topics outside the context, politely redirect to relevant topics
|
||||
|
||||
{custom_instructions}
|
||||
|
||||
Context from knowledge base:
|
||||
{context}
|
||||
"""
|
||||
|
||||
|
||||
class RAGEngine:
|
||||
def __init__(self):
|
||||
self.embedding_svc = embedding_service
|
||||
self.vector_svc = vector_store
|
||||
self.llm_svc = llm_service
|
||||
|
||||
async def process_query(
|
||||
self,
|
||||
query: str,
|
||||
collection_name: str,
|
||||
chatbot_config: Dict[str, Any],
|
||||
conversation_history: List[Dict[str, str]] = None,
|
||||
language: str = "en",
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Full RAG pipeline: embed → retrieve → generate
|
||||
"""
|
||||
if conversation_history is None:
|
||||
conversation_history = []
|
||||
|
||||
# Step 1: Embed the query
|
||||
try:
|
||||
query_embedding = self.embedding_svc.embed_text(query)
|
||||
except Exception as e:
|
||||
logger.error(f"Embedding error: {e}")
|
||||
return {
|
||||
"response": "I'm having trouble processing your request. Please try again.",
|
||||
"sources": [],
|
||||
"tokens_used": 0,
|
||||
"model": chatbot_config.get("model", "unknown"),
|
||||
}
|
||||
|
||||
# Step 2: Retrieve relevant chunks
|
||||
retrieved = self.vector_svc.search(
|
||||
collection_name=collection_name,
|
||||
query_vector=query_embedding,
|
||||
limit=5,
|
||||
score_threshold=0.3,
|
||||
)
|
||||
|
||||
# Step 3: Build sources
|
||||
sources = []
|
||||
context_parts = []
|
||||
seen_texts = set()
|
||||
|
||||
for item in retrieved:
|
||||
payload = item.get("payload", {})
|
||||
text = payload.get("text", "")
|
||||
if text and text not in seen_texts:
|
||||
seen_texts.add(text)
|
||||
context_parts.append(text)
|
||||
sources.append(
|
||||
SourceDocument(
|
||||
document_name=payload.get("file_name", "Document"),
|
||||
chunk_text=text[:200] + "..." if len(text) > 200 else text,
|
||||
score=item.get("score", 0.0),
|
||||
page_number=payload.get("page_number"),
|
||||
)
|
||||
)
|
||||
|
||||
context = "\n\n---\n\n".join(context_parts) if context_parts else "No relevant information found."
|
||||
|
||||
# Step 4: Build messages
|
||||
system_prompt = RAG_SYSTEM_PROMPT.format(
|
||||
company_name=chatbot_config.get("company_name", ""),
|
||||
custom_instructions=chatbot_config.get("system_prompt") or "",
|
||||
context=context,
|
||||
)
|
||||
|
||||
messages = [{"role": "system", "content": system_prompt}]
|
||||
|
||||
# Add conversation history (last 10 messages)
|
||||
for msg in conversation_history[-10:]:
|
||||
messages.append({"role": msg["role"], "content": msg["content"]})
|
||||
|
||||
# Add current query
|
||||
messages.append({"role": "user", "content": query})
|
||||
|
||||
# Step 5: Generate response
|
||||
model = chatbot_config.get("model", "accounts/fireworks/models/llama-v3p1-70b-instruct")
|
||||
try:
|
||||
result = await self.llm_svc.generate(
|
||||
messages=messages,
|
||||
model=model,
|
||||
max_tokens=chatbot_config.get("max_tokens", 1000),
|
||||
temperature=chatbot_config.get("temperature", 0.7),
|
||||
)
|
||||
return {
|
||||
"response": result["content"],
|
||||
"sources": sources,
|
||||
"tokens_used": result.get("tokens_used", 0),
|
||||
"model": result.get("model", model),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"LLM generation error: {e}")
|
||||
return {
|
||||
"response": "I'm having trouble generating a response. Please try again later.",
|
||||
"sources": sources,
|
||||
"tokens_used": 0,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
|
||||
rag_engine = RAGEngine()
|
||||
153
app/services/vector_store.py
Normal file
153
app/services/vector_store.py
Normal file
@@ -0,0 +1,153 @@
|
||||
from qdrant_client import QdrantClient, models
|
||||
from qdrant_client.http.models import (
|
||||
Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
|
||||
)
|
||||
from app.config import settings
|
||||
from typing import List, Dict, Any, Optional
|
||||
import logging
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_qdrant_client: QdrantClient = None
|
||||
|
||||
|
||||
def get_qdrant_client() -> QdrantClient:
|
||||
global _qdrant_client
|
||||
if _qdrant_client is None:
|
||||
kwargs = {"url": settings.qdrant_url}
|
||||
if settings.qdrant_api_key:
|
||||
kwargs["api_key"] = settings.qdrant_api_key
|
||||
_qdrant_client = QdrantClient(**kwargs)
|
||||
return _qdrant_client
|
||||
|
||||
|
||||
class VectorStoreService:
|
||||
VECTOR_SIZE = 1536 # text-embedding-3-small
|
||||
|
||||
def __init__(self):
|
||||
self.client = get_qdrant_client()
|
||||
|
||||
def create_collection(self, collection_name: str) -> bool:
|
||||
"""Create a new collection for a chatbot"""
|
||||
try:
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=self.VECTOR_SIZE,
|
||||
distance=Distance.COSINE,
|
||||
),
|
||||
)
|
||||
logger.info(f"Created collection: {collection_name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
if "already exists" in str(e).lower():
|
||||
return True
|
||||
logger.error(f"Error creating collection {collection_name}: {e}")
|
||||
raise
|
||||
|
||||
def delete_collection(self, collection_name: str) -> bool:
|
||||
"""Delete a chatbot's collection"""
|
||||
try:
|
||||
self.client.delete_collection(collection_name=collection_name)
|
||||
logger.info(f"Deleted collection: {collection_name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting collection {collection_name}: {e}")
|
||||
return False
|
||||
|
||||
def collection_exists(self, collection_name: str) -> bool:
|
||||
try:
|
||||
self.client.get_collection(collection_name)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def upsert_vectors(
|
||||
self,
|
||||
collection_name: str,
|
||||
vectors: List[List[float]],
|
||||
payloads: List[Dict[str, Any]],
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
"""Upsert vectors into collection"""
|
||||
if ids is None:
|
||||
ids = [str(uuid.uuid4()) for _ in vectors]
|
||||
|
||||
points = [
|
||||
PointStruct(
|
||||
id=idx,
|
||||
vector=vector,
|
||||
payload=payload,
|
||||
)
|
||||
for idx, vector, payload in zip(ids, vectors, payloads)
|
||||
]
|
||||
|
||||
try:
|
||||
self.client.upsert(
|
||||
collection_name=collection_name,
|
||||
points=points,
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error upserting vectors: {e}")
|
||||
raise
|
||||
|
||||
def search(
|
||||
self,
|
||||
collection_name: str,
|
||||
query_vector: List[float],
|
||||
limit: int = 5,
|
||||
score_threshold: float = 0.3,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Search for similar vectors"""
|
||||
try:
|
||||
results = self.client.search(
|
||||
collection_name=collection_name,
|
||||
query_vector=query_vector,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
return [
|
||||
{
|
||||
"id": str(r.id),
|
||||
"score": r.score,
|
||||
"payload": r.payload,
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
except Exception as e:
|
||||
logger.error(f"Error searching vectors: {e}")
|
||||
return []
|
||||
|
||||
def delete_by_document_id(self, collection_name: str, document_id: str) -> bool:
|
||||
"""Delete all vectors for a document"""
|
||||
try:
|
||||
self.client.delete(
|
||||
collection_name=collection_name,
|
||||
points_selector=models.FilterSelector(
|
||||
filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="document_id",
|
||||
match=MatchValue(value=document_id),
|
||||
)
|
||||
]
|
||||
)
|
||||
),
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting document vectors: {e}")
|
||||
return False
|
||||
|
||||
def count_vectors(self, collection_name: str) -> int:
|
||||
"""Count vectors in a collection"""
|
||||
try:
|
||||
result = self.client.count(collection_name=collection_name)
|
||||
return result.count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
vector_store = VectorStoreService()
|
||||
Reference in New Issue
Block a user