Merge pull request #4 from getzep/local_embedding

add env switch for embedding. default to off

Makefile

@@ -6,10 +6,10 @@ LANGUAGE_MODEL := en_core_web_sm
 all: download-language-model format lint test
 
 run:
-	poetry run python main.py
+	ENABLE_EMBEDDINGS=true poetry run python main.py
 
 run-dev:
-	poetry run uvicorn main:app --reload --log-level debug --port 8080
+	ENABLE_EMBEDDINGS=true poetry run uvicorn main:app --reload --log-level debug --port 8080
 
 docker-build:
 	DOCKER_BUILDKIT=1 docker build -t $(CONTAINER_NAME) .
@@ -36,7 +36,7 @@ lint:
 	poetry run ruff .
 
 test:
-	poetry run pytest app/tests
+	ENABLE_EMBEDDINGS=true poetry run pytest app/tests
 
 loadtest:
-	poetry run locust -f load_test.py --headless --run-time 60 -u 50 --host http://0.0.0.0:8080
+	ENABLE_EMBEDDINGS=true poetry run locust -f load_test.py --headless --run-time 60 -u 50 --host http://0.0.0.0:8080

app/api.py

@@ -2,9 +2,11 @@
 # Licensed under the MIT License.
 # Heavy modified by Zep
 
+import os
+
 import spacy
 import srsly  # type: ignore
-from fastapi import Body, FastAPI, status
+from fastapi import Body, FastAPI, HTTPException, status
 from fastapi.responses import ORJSONResponse
 from starlette.responses import PlainTextResponse, RedirectResponse
 
@@ -13,6 +15,8 @@
 from app.entity_extractor import SpacyExtractor
 from app.entity_models import Request, Response
 
+ENABLE_EMBEDDINGS = os.getenv("ENABLE_EMBEDDINGS", "false").lower() == "true"
+
 app = FastAPI(
     title="zep-nlp-server",
     version="0.2",
@@ -24,7 +28,8 @@
 nlp = spacy.load("en_core_web_sm")
 extractor = SpacyExtractor(nlp)
 
-embedder = Embedder()
+if ENABLE_EMBEDDINGS:
+    embedder = Embedder()
 
 
 @app.get("/healthz", response_model=str, status_code=status.HTTP_200_OK)
@@ -47,4 +52,7 @@ async def extract_entities(body: Request = Body(..., example=example_request)):
 @app.post("/embeddings", response_class=ORJSONResponse)
 async def embed_collection(collection: Collection):
     """Embed a Collection of Documents."""
-    return ORJSONResponse(embedder.embed(collection))
+    if ENABLE_EMBEDDINGS:
+        return ORJSONResponse(embedder.embed(collection))
+    else:
+        raise HTTPException(status_code=400, detail="Embeddings not enabled.")