Merge branch 'upstash-chat-store-integration' of https://github.com/f…

…ahreddinozcan/llama_index into upstash-chat-store-integration
run-llama · Sep 27, 2024 · ee08624 · ee08624
2 parents 69d3957 + ec9b8a8
commit ee08624
Show file tree

Hide file tree

Showing 9 changed files with 19 additions and 14 deletions.
diff --git a/llama-index-integrations/llms/llama-index-llms-ollama/llama_index/llms/ollama/base.py b/llama-index-integrations/llms/llama-index-llms-ollama/llama_index/llms/ollama/base.py
@@ -99,6 +99,10 @@ class Ollama(FunctionCallingLLM):
         default=True,
         description="Whether the model is a function calling model.",
     )
+    keep_alive: Optional[Union[float, str]] = Field(
+        default="5m",
+        description="controls how long the model will stay loaded into memory following the request(default: 5m)",
+    )
 
     _client: Optional[Client] = PrivateAttr()
     _async_client: Optional[AsyncClient] = PrivateAttr()
@@ -116,6 +120,7 @@ def __init__(
         client: Optional[Client] = None,
         async_client: Optional[AsyncClient] = None,
         is_function_calling_model: bool = True,
+        keep_alive: Optional[Union[float, str]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -128,6 +133,7 @@ def __init__(
             json_mode=json_mode,
             additional_kwargs=additional_kwargs,
             is_function_calling_model=is_function_calling_model,
+            keep_alive=keep_alive,
             **kwargs,
         )
 
@@ -279,6 +285,7 @@ def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
             format="json" if self.json_mode else "",
             tools=tools,
             options=self._model_kwargs,
+            keep_alive=self.keep_alive,
         )
 
         tool_calls = response["message"].get("tool_calls", [])
@@ -311,6 +318,7 @@ def gen() -> ChatResponseGen:
                 format="json" if self.json_mode else "",
                 tools=tools,
                 options=self._model_kwargs,
+                keep_alive=self.keep_alive,
             )
 
             response_txt = ""
@@ -354,6 +362,7 @@ async def gen() -> ChatResponseAsyncGen:
                 format="json" if self.json_mode else "",
                 tools=tools,
                 options=self._model_kwargs,
+                keep_alive=self.keep_alive,
             )
 
             response_txt = ""
@@ -396,6 +405,7 @@ async def achat(
             format="json" if self.json_mode else "",
             tools=tools,
             options=self._model_kwargs,
+            keep_alive=self.keep_alive,
         )
 
         tool_calls = response["message"].get("tool_calls", [])

diff --git a/llama-index-integrations/llms/llama-index-llms-ollama/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-ollama/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-ollama"
 readme = "README.md"
-version = "0.3.2"
+version = "0.3.3"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"

diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/BUILD
@@ -1,3 +1,4 @@
 poetry_requirements(
     name="poetry",
+    module_mapping={"pyjwt": ["jwt"]},
 )
diff --git a/...ctor_stores/llama-index-vector-stores-deeplake/llama_index/vector_stores/deeplake/base.py b/...ctor_stores/llama-index-vector-stores-deeplake/llama_index/vector_stores/deeplake/base.py
@@ -22,12 +22,7 @@
     node_to_metadata_dict,
 )
 
-try:
-    from deeplake.core.vectorstore.deeplake_vectorstore import VectorStore
-
-    DEEPLAKE_INSTALLED = True
-except ImportError:
-    DEEPLAKE_INSTALLED = False
+from deeplake.core.vectorstore.deeplake_vectorstore import VectorStore
 
 logger = logging.getLogger(__name__)
 
@@ -112,12 +107,6 @@ def __init__(
             num_workers=ingestion_num_workers,
         )
 
-        if not DEEPLAKE_INSTALLED:
-            raise ImportError(
-                "Could not import deeplake python package. "
-                "Please install it with `pip install deeplake`."
-            )
-
         self._vectorstore = VectorStore(
             path=dataset_path,
             ingestion_batch_size=ingestion_batch_size,

diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/pyproject.toml
@@ -27,12 +27,13 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-deeplake"
 readme = "README.md"
-version = "0.2.0"
+version = "0.2.1"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
 deeplake = ">=3.9.12"
 llama-index-core = "^0.11.0"
+pyjwt = "*"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "8.10.0"

diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/BUILD
@@ -1,3 +1,4 @@
 python_tests(
     interpreter_constraints=["==3.9.*", "==3.10.*"],
+    dependencies=["llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake:poetry#pyjwt"],
 )
diff --git a/...ons/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py b/...ons/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py
@@ -1,4 +1,5 @@
 import pytest
+import jwt  # noqa
 
 from llama_index.core import Document
 from llama_index.core.vector_stores.types import (

diff --git a/llama-index-packs/llama-index-packs-deeplake-deepmemory-retriever/tests/BUILD b/llama-index-packs/llama-index-packs-deeplake-deepmemory-retriever/tests/BUILD
@@ -1,3 +1,4 @@
 python_tests(
     interpreter_constraints=["==3.9.*", "==3.10.*"],
+    dependencies=["llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake:poetry#pyjwt"],
 )
diff --git a/llama-index-packs/llama-index-packs-deeplake-multimodal-retrieval/tests/BUILD b/llama-index-packs/llama-index-packs-deeplake-multimodal-retrieval/tests/BUILD
@@ -1,3 +1,4 @@
 python_tests(
     interpreter_constraints=["==3.9.*", "==3.10.*"],
+    dependencies=["llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake:poetry#pyjwt"]
 )