Upgrade redisvl (#32)

* update redisvl

* updates

* cleanup

* make an autouse fixture to clean up index client

* Clear logging handlers to prevent I/O errors in tests

* update to use lifespan manager from asgi

* remove fixtures

* update

---------

Co-authored-by: Andrew Brookins <[email protected]>

Author: tylerhutcherson

.gitignore

@@ -14,4 +14,5 @@ new_backend/arxivsearch/templates/
 .coverage*
 coverage.*
 htmlcov/
-legacy-data/
+legacy-data/
+.python-version

Dockerfile

@@ -35,7 +35,7 @@ RUN mkdir -p /app/backend
 COPY ./backend/poetry.lock ./backend/pyproject.toml ./backend/
 
 WORKDIR /app/backend
-RUN poetry install --all-extras --no-interaction --no-root
+RUN poetry install --all-extras --no-interaction
 
 COPY ./backend/ .
 

Makefile

@@ -1,4 +1,4 @@
-.PHONY: build
+.PHONY: deploy
 
-build:
-	docker compose -f docker-local-redis.yml up
+deploy:
+	docker compose -f docker-local-redis.yml up

README.md

@@ -3,6 +3,7 @@
     <a href="https://github.com/redis-developer/redis-arxiv-search"><img src="https://redis.io/wp-content/uploads/2024/04/Logotype.svg?raw=true" width="30%"><img></a>
     <br />
     <br />
+    <h1>🔎 arXiv Search API</h1>
 <div display="inline-block">
     <a href="https://docsearch.redisvl.com"><b>Hosted Demo</b></a>&nbsp;&nbsp;&nbsp;
     <a href="https://github.com/redis-developer/redis-arxiv-search"><b>Code</b></a>&nbsp;&nbsp;&nbsp;
@@ -14,7 +15,7 @@
     <br />
 </div>
 
-# 🔎 Redis arXiv Search
+
 *This repository is the official codebase for the arxiv paper search app hosted at: **https://docsearch.redisvl.com***
 
 
@@ -111,17 +112,17 @@ Embeddings represent the semantic properies of the raw text and enable vector si
     - Add your `OPENAI_API_KEY` to the `.env` file. **Need one?** [Get an API key](https://platform.openai.com)
     - Add you `COHERE_API_KEY` to the `.env` file. **Need one?** [Get an API key](https://cohere.ai)
 
-### Redis Stack Docker (Local) with make
+### Run locally with Redis 8 CE
 ```bash
-make build
+make deploy
 ```
 
 
 ## Customizing (optional)
 
 ### Run local redis with Docker
 ```bash
-docker run -d --name redis -p 6379:6379 -p 8001:8001 redis/redis-stack:latest
+docker run -d --name redis -p 6379:6379 -p 8001:8001 redis:8.0-M03
 ```
 
 ### FastApi with poetry

backend/arxivsearch/api/routes/papers.py

@@ -2,18 +2,19 @@
 import logging
 
 import numpy as np
+from fastapi import APIRouter, Depends, Query
+from redisvl.index import AsyncSearchIndex
+from redisvl.query import CountQuery, FilterQuery, VectorQuery
+
 from arxivsearch import config
-from arxivsearch.db import redis_helpers
+from arxivsearch.db import utils
 from arxivsearch.schema.models import (
     PaperSimilarityRequest,
     SearchResponse,
     UserTextSimilarityRequest,
     VectorSearchResponse,
 )
 from arxivsearch.utils.embeddings import Embeddings
-from fastapi import APIRouter, Depends, Query
-from redisvl.index import AsyncSearchIndex
-from redisvl.query import CountQuery, FilterQuery, VectorQuery
 
 logger = logging.getLogger(__name__)
 
@@ -26,7 +27,7 @@
 
 @router.get("/", response_model=SearchResponse)
 async def get_papers(
-    index: AsyncSearchIndex = Depends(redis_helpers.get_async_index),
+    index: AsyncSearchIndex = Depends(utils.get_async_index),
     limit: int = Query(default=20, description="Maximum number of papers to return."),
     skip: int = Query(
         default=0, description="Number of papers to skip for pagination."
@@ -53,9 +54,8 @@ async def get_papers(
     Returns:
         SearchResponse: Pydantic model containing papers and total count.
     """
-
     # Build queries
-    filter_expression = redis_helpers.build_filter_expression(
+    filter_expression = utils.build_filter_expression(
         years.split(","), categories.split(",")
     )
     filter_query = FilterQuery(return_fields=[], filter_expression=filter_expression)
@@ -72,7 +72,7 @@ async def get_papers(
 @router.post("/vector_search/by_paper", response_model=VectorSearchResponse)
 async def find_papers_by_paper(
     similarity_request: PaperSimilarityRequest,
-    index: AsyncSearchIndex = Depends(redis_helpers.get_async_index),
+    index: AsyncSearchIndex = Depends(utils.get_async_index),
 ):
     """
     Find and return papers similar to a given paper based on vector
@@ -93,7 +93,7 @@ async def find_papers_by_paper(
         paper[similarity_request.provider.value], dtype=np.float32
     )
     # Build filter expression
-    filter_expression = redis_helpers.build_filter_expression(
+    filter_expression = utils.build_filter_expression(
         similarity_request.years, similarity_request.categories
     )
     # Create queries
@@ -115,7 +115,7 @@ async def find_papers_by_paper(
 @router.post("/vector_search/by_text", response_model=VectorSearchResponse)
 async def find_papers_by_text(
     similarity_request: UserTextSimilarityRequest,
-    index: AsyncSearchIndex = Depends(redis_helpers.get_async_index),
+    index: AsyncSearchIndex = Depends(utils.get_async_index),
 ):
     """
     Find and return papers similar to user-provided text based on
@@ -131,7 +131,7 @@ async def find_papers_by_text(
     """
 
     # Build filter expression
-    filter_expression = redis_helpers.build_filter_expression(
+    filter_expression = utils.build_filter_expression(
         similarity_request.years, similarity_request.categories
     )
     # Check available paper count and create vector from user text

backend/arxivsearch/db/load.py

@@ -9,7 +9,7 @@
 from redisvl.index import AsyncSearchIndex
 
 from arxivsearch import config
-from arxivsearch.db import redis_helpers
+from arxivsearch.db.utils import get_async_index, get_schema
 from arxivsearch.schema.models import Provider
 
 logger = logging.getLogger(__name__)
@@ -74,8 +74,7 @@ async def preprocess_paper(paper: dict) -> dict:
 
 async def load_data():
     # Load schema specs and create index in Redis
-    index = AsyncSearchIndex(redis_helpers.schema)
-    await index.set_client(redis_helpers.client)
+    index = await get_async_index()
 
     # Load dataset and create index
     try:

backend/arxivsearch/db/index.yaml backend/arxivsearch/db/schema/index.yamlbackend/arxivsearch/db/index.yaml renamed to backend/arxivsearch/db/schema/index.yaml

@@ -2,42 +2,30 @@
 import os
 from typing import List
 
-from arxivsearch import config
-from redis.asyncio import Redis
-from redisvl.index import AsyncSearchIndex, SearchIndex
+from redisvl.index import AsyncSearchIndex
 from redisvl.query.filter import FilterExpression, Tag
 from redisvl.schema import IndexSchema
 
-logger = logging.getLogger(__name__)
-
-
-dir_path = os.path.dirname(os.path.realpath(__file__))
-schema_path = os.path.join(dir_path, "index.yaml")
-schema = IndexSchema.from_yaml(schema_path)
-client = Redis.from_url(config.REDIS_URL)
-global_index = None
-
+from arxivsearch import config
 
-def get_schema():
-    return IndexSchema.from_yaml(schema_path)
+logger = logging.getLogger(__name__)
 
 
-def get_test_index():
-    index = SearchIndex.from_yaml(schema_path)
-    index.connect(redis_url=config.REDIS_URL)
+# global search index
+_global_index = None
 
-    if not index.exists():
-        index.create(overwrite=True)
 
-    return index
+def get_schema() -> IndexSchema:
+    dir_path = os.path.dirname(os.path.realpath(__file__)) + "/schema"
+    file_path = os.path.join(dir_path, "index.yaml")
+    return IndexSchema.from_yaml(file_path)
 
 
 async def get_async_index():
-    global global_index
-    if not global_index:
-        global_index = AsyncSearchIndex.from_yaml(schema_path)
-        await global_index.set_client(client)
-    yield global_index
+    global _global_index
+    if not _global_index:
+        _global_index = AsyncSearchIndex(get_schema(), redis_url=config.REDIS_URL)
+    return _global_index
 
 
 def build_filter_expression(

backend/arxivsearch/db/redis_helpers.py backend/arxivsearch/db/utils.pybackend/arxivsearch/db/redis_helpers.py renamed to backend/arxivsearch/db/utils.py

@@ -1,20 +1,34 @@
 import logging
+from contextlib import asynccontextmanager
 from pathlib import Path
 
 import uvicorn
-from arxivsearch import config
-from arxivsearch.api.main import api_router
-from arxivsearch.spa import SinglePageApplication
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from starlette.middleware.cors import CORSMiddleware
 
+from arxivsearch import config
+from arxivsearch.api.main import api_router
+from arxivsearch.db.utils import get_async_index
+from arxivsearch.spa import SinglePageApplication
+
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    index = await get_async_index()
+    async with index:
+        yield
+
+
 app = FastAPI(
-    title=config.PROJECT_NAME, docs_url=config.API_DOCS, openapi_url=config.OPENAPI_DOCS
+    title=config.PROJECT_NAME,
+    docs_url=config.API_DOCS,
+    openapi_url=config.OPENAPI_DOCS,
+    lifespan=lifespan,
 )
 
 app.add_middleware(

backend/arxivsearch/main.py

@@ -1,4 +1,5 @@
 from enum import Enum
+
 from pydantic import BaseModel