Skip to content

Commit

Permalink
⚡️ Speed up method AstraDBVectorStoreComponent.get_database_object
Browse files Browse the repository at this point in the history
…by 1,269% in PR #6236 (`LFOSS-492`)

To optimize the provided code for better performance, we can implement a few strategies.

Below is the refactored code with these performance optimizations.

### Key Changes.
  • Loading branch information
codeflash-ai[bot] authored Feb 14, 2025
1 parent 97b9380 commit f55a73f
Showing 1 changed file with 35 additions and 13 deletions.
48 changes: 35 additions & 13 deletions src/backend/base/langflow/components/vectorstores/astradb.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import functools
from collections import defaultdict
from dataclasses import asdict, dataclass, field

Expand Down Expand Up @@ -433,30 +434,32 @@ def get_api_endpoint_static(
return db.get("api_endpoint")

def get_api_endpoint(self):
return self.get_api_endpoint_static(
token=self.token,
environment=self.environment,
api_endpoint=self.api_endpoint,
database_name=self.database_name,
)
if not self.api_endpoint_cache:
self.api_endpoint_cache = self.get_api_endpoint_static_cached(
token=self.token,
environment=self.environment,
api_endpoint=self.api_endpoint,
database_name=self.database_name,
)
return self.api_endpoint_cache

def get_keyspace(self):
keyspace = self.keyspace

if keyspace:
return keyspace.strip()

if self.keyspace:
return self.strip_keyspace(self.keyspace)
return None

def get_database_object(self, api_endpoint: str | None = None):
if self.database_object_cache:
return self.database_object_cache

try:
client = DataAPIClient(token=self.token, environment=self.environment)

return client.get_database(
self.database_object_cache = client.get_database(
api_endpoint=api_endpoint or self.get_api_endpoint(),
token=self.token,
keyspace=self.get_keyspace(),
)
return self.database_object_cache

Check failure on line 462 in src/backend/base/langflow/components/vectorstores/astradb.py

View workflow job for this annotation

GitHub Actions / Ruff Style Check (3.12)

Ruff (TRY300)

src/backend/base/langflow/components/vectorstores/astradb.py:462:13: TRY300 Consider moving this statement to an `else` block
except Exception as e:
msg = f"Error fetching database object: {e}"
raise ValueError(msg) from e
Expand Down Expand Up @@ -946,3 +949,22 @@ def get_retriever_kwargs(self):
"search_type": self._map_search_type(),
"search_kwargs": search_args,
}

def __init__(self, **kwargs):
super().__init__(**kwargs)
self.api_endpoint_cache = None
self.database_object_cache = None

@functools.lru_cache(maxsize=128)

Check failure on line 958 in src/backend/base/langflow/components/vectorstores/astradb.py

View workflow job for this annotation

GitHub Actions / Ruff Style Check (3.12)

Ruff (B019)

src/backend/base/langflow/components/vectorstores/astradb.py:958:5: B019 Use of `functools.lru_cache` or `functools.cache` on methods can lead to memory leaks
def get_api_endpoint_static_cached(self, token, environment, api_endpoint, database_name):
return self.get_api_endpoint_static(
token=token,
environment=environment,
api_endpoint=api_endpoint,
database_name=database_name,
)

@staticmethod
@functools.lru_cache(maxsize=128)
def strip_keyspace(keyspace):
return keyspace.strip()

0 comments on commit f55a73f

Please # to comment.