chore: import upstream snapshot with attribution

tests/llama_index/sample_code/basic_chat_engine.py

@@ -0,0 +1,16 @@
+"""
+Sample code to define a chat engine and save it with model-from-code logging.
+
+Ref: https://qdrant.tech/documentation/quickstart/
+"""
+
+from llama_index.core import Document, VectorStoreIndex
+from llama_index.core.chat_engine.types import ChatMode
+
+import mlflow
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+# Setting SIMPLE chat mode will create a SimpleChatEngine instance
+chat_engine = index.as_chat_engine(chat_mode=ChatMode.SIMPLE)
+
+mlflow.models.set_model(chat_engine)

tests/llama_index/sample_code/basic_retriever.py

@@ -0,0 +1,8 @@
+from llama_index.core import Document, VectorStoreIndex
+
+import mlflow
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+retriever = index.as_retriever()
+
+mlflow.models.set_model(retriever)

tests/llama_index/sample_code/basic_vector_store.py

@@ -0,0 +1,7 @@
+from llama_index.core import Document, VectorStoreIndex
+
+import mlflow
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+
+mlflow.models.set_model(index)

tests/llama_index/sample_code/external_vector_store.py

@@ -0,0 +1,36 @@
+"""
+Sample code to define an index using an external vector store (Faiss) for model-from-code logging.
+
+Ref: https://qdrant.tech/documentation/quickstart/
+"""
+
+from llama_index.core import VectorStoreIndex
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, PointStruct, VectorParams
+
+import mlflow
+
+# Run Qdrant with in-memory mode for testing purpose
+client = QdrantClient(location=":memory:")
+client.create_collection(
+    collection_name="test",
+    # 1536 is the size of OpenAI embeddings
+    vectors_config=VectorParams(size=1536, distance=Distance.DOT),
+)
+# dummy embeddings
+vec = [0.1] * 1536
+client.upsert(
+    collection_name="test",
+    wait=True,
+    points=[
+        PointStruct(id=1, vector=vec, payload={"text": "hi"}),
+        PointStruct(id=1, vector=vec, payload={"text": "hola"}),
+    ],
+)
+
+vector_store = QdrantVectorStore(client=client, collection_name="test")
+index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+
+
+mlflow.models.set_model(index)

tests/llama_index/sample_code/model_config.yaml

@@ -0,0 +1,3 @@
+llm:
+  model_name: "gpt-4o-mini"
+  temperature: 0.7

tests/llama_index/sample_code/query_engine_with_reranker.py

@@ -0,0 +1,41 @@
+"""
+Sample code to define a query engine with post processors and save it with model-from-code logging.
+
+Ref: https://qdrant.tech/documentation/quickstart/
+"""
+
+from llama_index.core import Document, QueryBundle, VectorStoreIndex
+from llama_index.core.postprocessor import LLMRerank
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.schema import NodeWithScore
+
+import mlflow
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+
+# Postprocessor 1: Reranker
+reranker = LLMRerank(
+    choice_batch_size=5,
+    top_n=2,
+)
+
+
+# Postprocessor 2: custom postprocessor
+class CustomNodePostprocessor(BaseNodePostprocessor):
+    call_count: int = 0
+
+    def _postprocess_nodes(
+        self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None
+    ) -> list[NodeWithScore]:
+        # subtracts 1 from the score
+        self.call_count += 1
+        return nodes
+
+
+query_engine = index.as_query_engine(
+    similarity_top_k=5,
+    node_postprocessors=[reranker, CustomNodePostprocessor()],
+)
+
+
+mlflow.models.set_model(query_engine)

tests/llama_index/sample_code/simple_workflow.py

@@ -0,0 +1,37 @@
+from llama_index.core.workflow import (
+    Event,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+from llama_index.llms.openai import OpenAI
+
+import mlflow
+
+
+class JokeEvent(Event):
+    joke: str
+
+
+class JokeFlow(Workflow):
+    llm = OpenAI()
+
+    @step
+    async def generate_joke(self, ev: StartEvent) -> JokeEvent:
+        topic = ev.topic
+        prompt = f"Write your best joke about {topic}."
+        response = await self.llm.acomplete(prompt)
+        return JokeEvent(joke=str(response))
+
+    @step
+    async def critique_joke(self, ev: JokeEvent) -> StopEvent:
+        joke = ev.joke
+
+        prompt = f"Give a thorough analysis and critique of the following joke: {joke}"
+        response = await self.llm.acomplete(prompt)
+        return StopEvent(result=str(response))
+
+
+w = JokeFlow(timeout=10, verbose=False)
+mlflow.models.set_model(w)

tests/llama_index/sample_code/with_model_config.py

@@ -0,0 +1,20 @@
+"""
+Sample code to define a chat engine and save it with model config (dictionary).
+"""
+
+from llama_index.core import Document, VectorStoreIndex
+from llama_index.core.chat_engine.types import ChatMode
+from llama_index.llms.openai import OpenAI
+
+import mlflow
+
+model_config = mlflow.models.ModelConfig()
+model_name = model_config.get("model_name")
+temperature = model_config.get("temperature")
+llm = OpenAI(model=model_name, temperature=temperature)
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+# Setting SIMPLE chat mode will create a SimpleChatEngine instance
+chat_engine = index.as_chat_engine(llm=llm, chat_mode=ChatMode.SIMPLE)
+
+mlflow.models.set_model(chat_engine)

tests/llama_index/sample_code/with_model_config_yaml_file.py

@@ -0,0 +1,21 @@
+"""
+Sample code to define a chat engine and save it with model config (YAML file).
+"""
+
+from llama_index.core import Document, VectorStoreIndex
+from llama_index.core.chat_engine.types import ChatMode
+from llama_index.llms.openai import OpenAI
+
+import mlflow
+
+model_config = mlflow.models.ModelConfig(development_config="model_config.yaml")
+llm_config = model_config.get("llm")
+model_name = llm_config.get("model_name")
+temperature = llm_config.get("temperature")
+llm = OpenAI(model=model_name, temperature=temperature)
+
+index = VectorStoreIndex.from_documents(documents=[Document.example()])
+# Setting SIMPLE chat mode will create a SimpleChatEngine instance
+chat_engine = index.as_chat_engine(llm=llm, chat_mode=ChatMode.SIMPLE)
+
+mlflow.models.set_model(chat_engine)