Merge pull request #28 from epuerta9/streaming

Streaming

app.py

@@ -0,0 +1,72 @@
+from kitchenai.contrib.kitchenai_sdk.kitchenai import KitchenAIApp
+from kitchenai.contrib.kitchenai_sdk.api import QuerySchema, EmbedSchema
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.llms.openai import OpenAI
+import os 
+import chromadb
+from llama_index.llms.openai import OpenAI
+from llama_index.core.node_parser import TokenTextSplitter
+from llama_index.core.extractors import (
+    TitleExtractor,
+    QuestionsAnsweredExtractor)
+from llama_index.core import Document
+from kitchenai.contrib.kitchenai_sdk.storage.llama_parser import Parser
+
+kitchen = KitchenAIApp()
+
+chroma_client = chromadb.PersistentClient(path="chroma_db")
+chroma_collection = chroma_client.get_or_create_collection("quickstart")
+llm = OpenAI(model="gpt-4")
+chroma_collection_second_collection = chroma_client.get_or_create_collection("second_collection")
+
+@kitchen.storage("simple-vector")
+def simple_vector(dir: str, metadata: dict = {}, *args, **kwargs):
+    parser = Parser(api_key=os.environ.get("LLAMA_CLOUD_API_KEY", None))
+    response = parser.load(dir, metadata=metadata, **kwargs)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    VectorStoreIndex.from_documents(
+        response["documents"], storage_context=storage_context, show_progress=True,
+            transformations=[TokenTextSplitter(), TitleExtractor(),QuestionsAnsweredExtractor()]
+    )
+    return {"response": len(response["documents"])}
+
+
+@kitchen.storage("simple-vector2")
+def simple_vector2(dir: str, metadata: dict = {}, *args, **kwargs):
+    parser = Parser(api_key=os.environ.get("LLAMA_CLOUD_API_KEY", None))
+    response = parser.load(dir, metadata=metadata, **kwargs)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection_second_collection)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    VectorStoreIndex.from_documents(
+        response["documents"], storage_context=storage_context, show_progress=True,
+            transformations=[TokenTextSplitter(), TitleExtractor(),QuestionsAnsweredExtractor()]
+    )
+    return {"response": len(response["documents"])}
+
+@kitchen.query("simple-query")
+def simple_query(data: QuerySchema):
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    index = VectorStoreIndex.from_vector_store(
+        vector_store,
+    )
+    query_engine = index.as_query_engine(chat_mode="best", llm=llm, verbose=True)
+    response = query_engine.query(data.query)
+    return {"response": response.response}
+
+@kitchen.stream("simple-stream")
+def simple_stream(data: QuerySchema):
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    index = VectorStoreIndex.from_vector_store(
+        vector_store,
+    )
+    query_engine = index.as_query_engine(streaming=True, llm=llm, similarity_top_k=1, verbose=True)
+    response = query_engine.query(data.query)
+
+    return response.response_gen
+
+@kitchen.query("non-ai")
+def non_ai(data: QuerySchema):
+    msg = "no AI is used in this function"
+    return {"response": msg}

docs/requirements.txt

@@ -1,7 +1,7 @@
 #
 # This file is autogenerated by hatch-pip-compile with Python 3.11
 #
-# [constraints] requirements.txt (SHA256: bef74194f8fd80762a1ef00469312dbc8243a3854822b60f9150002d8d517687)
+# [constraints] requirements.txt (SHA256: 18df5e1d5be84eeae244327525e84700b34a08f7506779624aa1f99fe146898d)
 #
 # - django-browser-reload
 # - django-debug-toolbar
@@ -36,6 +36,7 @@
 # - django-allauth[socialaccount]
 # - django-anymail[resend]
 # - django-compressor
+# - django-eventstream
 # - django-extensions
 # - django-health-check
 # - django-htmx
@@ -59,7 +60,7 @@
 # - heroicons[django]
 # - honcho
 # - jupyterlab
-# - kitchenai-python-sdk
+# - kitchenai-python-sdk>=1.3.0
 # - llama-index
 # - llama-index-vector-stores-chroma
 # - nbconvert
@@ -357,8 +358,10 @@ django==5.1.2
     #   django-compressor
     #   django-crispy-forms
     #   django-debug-toolbar
+    #   django-eventstream
     #   django-extensions
     #   django-fastdev
+    #   django-grip
     #   django-health-check
     #   django-htmx
     #   django-lifecycle
@@ -411,13 +414,21 @@ django-crispy-forms==2.3
     #   crispy-tailwind
 django-debug-toolbar==4.4.6
     # via hatch.envs.docs
+django-eventstream==5.3.1
+    # via
+    #   -c requirements.txt
+    #   hatch.envs.docs
 django-extensions==3.2.3
     # via
     #   -c requirements.txt
     #   hatch.envs.docs
     #   dj-notebook
 django-fastdev==1.12.0
     # via hatch.envs.docs
+django-grip==3.5.1
+    # via
+    #   -c requirements.txt
+    #   django-eventstream
 django-health-check==3.18.3
     # via
     #   -c requirements.txt
@@ -592,6 +603,11 @@ greenlet==3.1.1
     # via
     #   -c requirements.txt
     #   sqlalchemy
+gripcontrol==4.2.0
+    # via
+    #   -c requirements.txt
+    #   django-eventstream
+    #   django-grip
 grpcio==1.68.1
     # via
     #   -c requirements.txt
@@ -826,7 +842,7 @@ jupyterlab-widgets==3.0.13
     # via ipywidgets
 keyring==25.4.1
     # via hatch
-kitchenai-python-sdk==1.1.0
+kitchenai-python-sdk==1.3.0
     # via
     #   -c requirements.txt
     #   hatch.envs.docs
@@ -1230,6 +1246,11 @@ ptyprocess==0.7.0
     #   -c requirements.txt
     #   pexpect
     #   terminado
+pubcontrol==3.5.0
+    # via
+    #   -c requirements.txt
+    #   django-grip
+    #   gripcontrol
 pure-eval==0.2.3
     # via
     #   -c requirements.txt
@@ -1282,6 +1303,9 @@ pyjwt==2.9.0
     # via
     #   -c requirements.txt
     #   django-allauth
+    #   django-eventstream
+    #   gripcontrol
+    #   pubcontrol
 pypdf==5.1.0
     # via
     #   -c requirements.txt
@@ -1393,6 +1417,7 @@ requests==2.32.3
     #   llama-index-core
     #   llama-index-legacy
     #   posthog
+    #   pubcontrol
     #   requests-oauthlib
     #   sphinx
     #   tiktoken
@@ -1471,7 +1496,10 @@ six==1.16.0
     #   -c requirements.txt
     #   asttokens
     #   cssbeautifier
+    #   django-eventstream
+    #   django-grip
     #   django-pandas
+    #   gripcontrol
     #   jsbeautifier
     #   kubernetes
     #   posthog
@@ -1759,8 +1787,11 @@ websockets==14.1
     #   -c requirements.txt
     #   sphinx-autobuild
     #   uvicorn
-werkzeug==3.0.4
-    # via hatch.envs.docs
+werkzeug==3.1.3
+    # via
+    #   -c requirements.txt
+    #   hatch.envs.docs
+    #   django-grip
 wheel==0.44.0
     # via pip-tools
 whitenoise==6.7.0

kitchenai/api.py

@@ -1,16 +1,14 @@
 import logging
 
 from django.apps import apps
-from kitchenai.core.api import router as core_router
 from ninja import NinjaAPI
+from kitchenai.core.router import router as core_router
 
 logger = logging.getLogger(__name__)
 
-api = NinjaAPI(version="0.9.0")
+api = NinjaAPI(version="0.10.0")
 
 
-# Ensure `ready()` is called for all apps
 apps.get_app_configs()
 
-# Get the app's router from `MyAppConfig` and add it to the main API
-api.add_router("/core", core_router)  # Add app's router here
+api.add_router("/v1", core_router)  

kitchenai/cli/client.py → kitchenai/cli/client/client.py

@@ -245,6 +245,4 @@ def list_labels():
             )
             console.print(table)
         except Exception as e:
-            console.print(f"[red]Error fetching labels: {e}[/red]")
-
-
+            console.print(f"[red]Error fetching labels: {e}[/red]")