chore: import upstream snapshot with attribution

simple-rag-workflow/workflow.py

@@ -0,0 +1,98 @@
+import nest_asyncio
+from llama_index.llms.ollama import Ollama
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.settings import Settings
+from llama_index.core.workflow import Event, Context, Workflow, StartEvent, StopEvent, step
+from llama_index.core.schema import NodeWithScore
+from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
+from llama_index.core.response_synthesizers import CompactAndRefine
+
+# Apply nest_asyncio to allow nested event loops
+nest_asyncio.apply()
+
+class RetrieverEvent(Event):
+    """Result of running retrieval"""
+    nodes: list[NodeWithScore]
+
+class RAGWorkflow(Workflow):
+    def __init__(self, model_name="llama3.2", embedding_model="BAAI/bge-small-en-v1.5"):
+        super().__init__()
+        # Initialize LLM and embedding model
+        self.llm = Ollama(model=model_name)
+        self.embed_model = HuggingFaceEmbedding(model_name=embedding_model)
+        
+        # Configure global settings
+        Settings.llm = self.llm
+        Settings.embed_model = self.embed_model
+        
+        self.index = None
+
+    @step
+    async def ingest(self, ctx: Context, ev: StartEvent) -> StopEvent | None:
+        """Entry point to ingest documents from a directory."""
+        dirname = ev.get("dirname")
+        if not dirname:
+            return None
+
+        documents = SimpleDirectoryReader(dirname).load_data()
+        self.index = VectorStoreIndex.from_documents(documents=documents)
+        return StopEvent(result=self.index)
+
+    @step
+    async def retrieve(self, ctx: Context, ev: StartEvent) -> RetrieverEvent | None:
+        """Entry point for RAG retrieval."""
+        query = ev.get("query")
+        index = ev.get("index") or self.index
+
+        if not query:
+            return None
+
+        if index is None:
+            print("Index is empty, load some documents before querying!")
+            return None
+
+        retriever = index.as_retriever(similarity_top_k=2)
+        nodes = await retriever.aretrieve(query)
+        await ctx.set("query", query)
+        return RetrieverEvent(nodes=nodes)
+
+    @step
+    async def synthesize(self, ctx: Context, ev: RetrieverEvent) -> StopEvent:
+        """Generate a response using retrieved nodes."""
+        summarizer = CompactAndRefine(streaming=True, verbose=True)
+        query = await ctx.get("query", default=None)
+        response = await summarizer.asynthesize(query, nodes=ev.nodes)
+        return StopEvent(result=response)
+
+    async def query(self, query_text: str):
+        """Helper method to perform a complete RAG query."""
+        if self.index is None:
+            raise ValueError("No documents have been ingested. Call ingest_documents first.")
+        
+        result = await self.run(query=query_text, index=self.index)
+        return result
+
+    async def ingest_documents(self, directory: str):
+        """Helper method to ingest documents."""
+        result = await self.run(dirname=directory)
+        self.index = result
+        return result
+
+# Example usage
+async def main():
+    # Initialize the workflow
+    workflow = RAGWorkflow()
+    
+    # Ingest documents
+    await workflow.ingest_documents("data")
+    
+    # Perform a query
+    result = await workflow.query("How was DeepSeekR1 trained?")
+    
+    # Print the response
+    async for chunk in result.async_response_gen():
+        print(chunk, end="", flush=True)
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())