Spaces:
Sleeping
Sleeping
File size: 6,587 Bytes
e6bfb30 35b8c1f 479a5d1 35b8c1f e6bfb30 35b8c1f 72db647 27f97be e6bfb30 35b8c1f e6bfb30 35b8c1f e6bfb30 6067005 e6bfb30 6067005 72db647 35b8c1f e6bfb30 4316268 35b8c1f e6bfb30 479a5d1 e6bfb30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
import json
import os
from typing import List, Dict, Any, Optional
from pydantic import BaseModel
import uvicorn
from fastapi import FastAPI, HTTPException
from pinecone import Pinecone , ServerlessSpec
import numpy as np
from openai import OpenAI
# Load environment variables
from dotenv import load_dotenv
load_dotenv()
# Get API keys from environment variables
PINECONE_API_KEY = os.getenv('PINECONE_API_KEY')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if not PINECONE_API_KEY:
raise ValueError("PINECONE_API_KEY environment variable not set")
if not OPENAI_API_KEY:
raise ValueError("OPENAI_API_KEY environment variable not set")
# Create FastAPI app
app = FastAPI(title="E-Bikes Semantic Search API",
description="API for finding similar e-bikes based on semantic search",
version="1.0.0")
def build_filter(pt: Optional[str], cat: Optional[str]) -> dict | None:
filt = {}
if pt:
filt["type"] = pt # shorthand $eq
if cat:
filt["category"] = cat
return filt or None
# Request and response models
class SearchRequest(BaseModel):
description: str
top_k: int = 3
product_type: str
category : str
class BikeMatch(BaseModel):
id: str
name: str
type: str
description: str
score: float
class SearchResponse(BaseModel):
matches: List[BikeMatch]
# Initialize OpenAI client
openai_client = OpenAI(api_key=OPENAI_API_KEY)
# Define the embedding model using OpenAI
class OpenAIEmbedder:
def __init__(self, model_name="text-embedding-3-small"):
self.model_name = model_name
self.client = openai_client
self.embedding_dimension = 1536 # Dimension of text-embedding-3-small
def encode(self, texts):
if isinstance(texts, str):
texts = [texts]
# Get embeddings from OpenAI
response = self.client.embeddings.create(
input=texts,
model=self.model_name
)
# Extract embeddings from response
embeddings = [item.embedding for item in response.data]
return np.array(embeddings)
# Initialize Pinecone client
def initialize_pinecone():
pc = Pinecone(api_key=PINECONE_API_KEY)
# Define index name
index_name = "ebikes-search"
# Check if index already exists
existing_indexes = pc.list_indexes().names()
if index_name not in existing_indexes:
# Create index with 1536 dimensions (matches text-embedding-3-small)
pc.create_index(
name=index_name,
dimension=1536,
metric="cosine",
spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
print(f"Created new index: {index_name}")
# Connect to the index
try:
index = pc.Index(index_name)
return index
except Exception as e:
print(f"Error connecting to Pinecone index: {e}")
raise
# Load the e-bikes data
def load_ebikes_data(file_path="data.json"):
try:
with open(file_path, 'r') as f:
data = json.load(f)
return data.get('pogo-cycles-data', [])
except Exception as e:
print(f"Error loading e-bikes data: {e}")
return []
# Create embeddings and upload to Pinecone
def create_and_upload_embeddings(ebikes_data, encoder, pinecone_index):
# Prepare data for indexing
ids = []
descriptions = []
metadata = []
for bike in ebikes_data:
ids.append(bike['id'])
descriptions.append(bike['description'])
metadata.append({
'id': bike['id'],
'name': bike['name'],
'type': bike['product_type'],
'description': bike['description'],
'category': bike['category']
})
# Create embeddings
embeddings = encoder.encode(descriptions)
# Prepare vectors for Pinecone
vectors_to_upsert = []
for i in range(len(ids)):
vector = {
'id': ids[i],
'values': embeddings[i].tolist(),
'metadata': metadata[i]
}
vectors_to_upsert.append(vector)
# Upsert vectors to Pinecone
pinecone_index.upsert(vectors=vectors_to_upsert)
print(f"Uploaded {len(vectors_to_upsert)} embeddings to Pinecone")
# Global variables for model and Pinecone index
encoder = None
pinecone_index = None
# Initialize data at startup
@app.on_event("startup")
async def startup_event():
global encoder, pinecone_index
print("Initializing OpenAI embedder...")
encoder = OpenAIEmbedder()
print("Connecting to Pinecone...")
pinecone_index = initialize_pinecone()
print("Loading e-bikes data...")
ebikes_data = load_ebikes_data("data.json")
if not ebikes_data:
print("No e-bikes data found, skipping embedding creation")
return
print("Creating and uploading embeddings...")
create_and_upload_embeddings(ebikes_data, encoder, pinecone_index)
print("API startup completed successfully!")
@app.get("/health")
async def health_check():
"""Health check endpoint"""
return {"status": "healthy"}
@app.post("/search", response_model=SearchResponse)
async def search_ebikes(request:SearchRequest):
"""
Search for e-bikes similar to the provided description
This endpoint uses semantic search to find e-bikes that match the user's description.
"""
try:
# Create embedding for the query
query_embedding = encoder.encode(request.description)[0]
filter_payload = build_filter(request.product_type, request.category)
# Query Pinecone
results = pinecone_index.query(
vector=query_embedding.tolist(),
top_k=3,
include_metadata=True,
filter=filter_payload
)
print("results",results)
# Parse results
matches = []
for match in results.matches:
bike_match = BikeMatch(
id=match.metadata.get('id'),
name=match.metadata.get('name'),
type=match.metadata.get('type'),
description=match.metadata.get('description'),
score=float(match.score)
)
matches.append(bike_match)
return SearchResponse(matches=matches)
except Exception as e:
print(f"Error during search: {e}")
raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True) |