rohitkshirsagar19 commited on
Commit
71a0948
·
verified ·
1 Parent(s): a72b521

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +123 -0
app.py ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import uvicorn
2
+ from fastapi import FastAPI, HTTPException, Depends
3
+ from fastapi.middleware.cors import CORSMiddleware
4
+ from pydantic import BaseModel
5
+ from sentence_transformers import SentenceTransformer
6
+ from pinecone import Pinecone, ServerlessSpec
7
+ import uuid
8
+ import os
9
+ from contextlib import asynccontextmanager
10
+
11
+ # --- Environment Setup ---
12
+ # It's best practice to get sensitive keys from environment variables
13
+ # We will set these up in Hugging Face Spaces Secrets
14
+ PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
15
+ PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "memoria-index")
16
+
17
+ # --- Global objects ---
18
+ # We load these once at startup to save time and memory
19
+ model = None
20
+ pc = None
21
+ index = None
22
+
23
+ @asynccontextmanager
24
+ async def lifespan(app: FastAPI):
25
+ """
26
+ Handles startup and shutdown events for the FastAPI app.
27
+ Loads the model and connects to Pinecone on startup.
28
+ """
29
+ global model, pc, index
30
+ print("Application startup...")
31
+
32
+ if not PINECONE_API_KEY:
33
+ raise ValueError("PINECONE_API_KEY environment variable not set.")
34
+
35
+ # 1. Load the AI Model
36
+ print("Loading lightweight sentence transformer model...")
37
+ model = SentenceTransformer('sentence-transformers/paraphrase-albert-small-v2')
38
+ print("Model loaded.")
39
+
40
+ # 2. Connect to Pinecone
41
+ print("Connecting to Pinecone...")
42
+ pc = Pinecone(api_key=PINECONE_API_KEY)
43
+
44
+ # 3. Get or create the Pinecone index
45
+ if PINECONE_INDEX_NAME not in pc.list_indexes().names():
46
+ print(f"Creating new Pinecone index: {PINECONE_INDEX_NAME}")
47
+ pc.create_index(
48
+ name=PINECONE_INDEX_NAME,
49
+ dimension=model.get_sentence_embedding_dimension(),
50
+ metric="cosine", # Cosine similarity is great for sentence vectors
51
+ spec=ServerlessSpec(cloud="aws", region="us-east-1")
52
+ )
53
+ index = pc.Index(PINECONE_INDEX_NAME)
54
+ print("Pinecone setup complete.")
55
+ yield
56
+ # Cleanup logic can go here if needed on shutdown
57
+ print("Application shutdown.")
58
+
59
+ # --- Pydantic Models ---
60
+ class Memory(BaseModel):
61
+ content: str
62
+
63
+ class SearchQuery(BaseModel):
64
+ query: str
65
+
66
+ # --- FastAPI App ---
67
+ app = FastAPI(
68
+ title="Memoria API",
69
+ description="API for storing and retrieving memories.",
70
+ version="1.0.0",
71
+ lifespan=lifespan # Use the lifespan context manager
72
+ )
73
+
74
+ app.add_middleware(
75
+ CORSMiddleware,
76
+ allow_origins=["*"], # Allow all origins for simplicity
77
+ allow_credentials=True,
78
+ allow_methods=["*"],
79
+ allow_headers=["*"],
80
+ )
81
+
82
+ # --- API Endpoints ---
83
+ @app.get("/")
84
+ def read_root():
85
+ return {"status": "ok", "message": "Welcome to the Memoria API!"}
86
+
87
+ @app.post("/save_memory")
88
+ def save_memory(memory: Memory):
89
+ try:
90
+ embedding = model.encode(memory.content).tolist()
91
+ memory_id = str(uuid.uuid4())
92
+
93
+ # Upsert (update or insert) the vector into Pinecone
94
+ index.upsert(vectors=[{"id": memory_id, "values": embedding, "metadata": {"text": memory.content}}])
95
+
96
+ print(f"Successfully saved memory with ID: {memory_id}")
97
+ return {"status": "success", "id": memory_id}
98
+ except Exception as e:
99
+ print(f"An error occurred during save: {e}")
100
+ raise HTTPException(status_code=500, detail=str(e))
101
+
102
+ @app.post("/search_memory")
103
+ def search_memory(search: SearchQuery):
104
+ try:
105
+ query_embedding = model.encode(search.query).tolist()
106
+
107
+ # Query Pinecone for the most similar vectors
108
+ results = index.query(vector=query_embedding, top_k=5, include_metadata=True)
109
+
110
+ # Extract the original text from the metadata
111
+ retrieved_documents = [match['metadata']['text'] for match in results['matches']]
112
+
113
+ print(f"Found {len(retrieved_documents)} results for query: '{search.query}'")
114
+ return {"status": "success", "results": retrieved_documents}
115
+ except Exception as e:
116
+ print(f"An error occurred during search: {e}")
117
+ raise HTTPException(status_code=500, detail=str(e))
118
+
119
+ if __name__ == "__main__":
120
+ uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)
121
+
122
+
123
+