first commit

.gitignore

@@ -0,0 +1,4 @@
+pw-env/
+.vscode/
+__pycache__/
+**/.env

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Pathway
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,80 @@
+
+<div align="center">
+  <img src="https://pathway.com/logo-light.svg" /><br /><br />
+</div>
+<p align="center">
+    <a href="https://github.com/pathwaycom/llm-app-pathway/blob/main/LICENSE">
+        <img src="https://img.shields.io/github/license/pathwaycom/llm-app-pathway?style=plastic" alt="Contributors"/></a>
+    <a href="https://github.com/pathwaycom/llm-app-pathway/graphs/contributors">
+        <img src="https://img.shields.io/github/contributors/pathwaycom/llm-app-pathway?style=plastic" alt="Contributors"/></a>
+    <a href="https://github.com/pathwaycom/llm-app-pathway/actions/workflows/install_package.yml">
+        <img src="https://img.shields.io/github/actions/workflow/status/pathwaycom/llm-app-pathway/install_package.yml?style=plastic" alt="Build" /></a> 
+        <img src="https://img.shields.io/badge/OS-Linux-green" alt="Linux"/>
+        <img src="https://img.shields.io/badge/OS-macOS-green" alt="macOS"/>
+      <br>
+    <a href="https://discord.gg/pathway">
+        <img src="https://img.shields.io/discord/1042405378304004156?logo=discord"
+            alt="chat on Discord"></a>
+    <a href="https://twitter.com/intent/follow?screen_name=pathway_com">
+        <img src="https://img.shields.io/twitter/follow/pathway_com?style=social&logo=twitter"
+            alt="follow on Twitter"></a>
+</p>
+
+This repository contains the code for our Pathway-based LLM App, which is designed to provide real-time responses to queries about the Pathway documentation.
+
+## Project Overview
+
+The app reads a corpus of documents stored in S3 or locally, preprocesses them, and builds a Pathway vector index. It then listens to user queries coming as HTTP REST requests. Each query uses the index to retrieve relevant documentation snippets and uses the OpenAI API/ Hugging Face to provide a response in natural language. The bot is reactive to changes in the corpus of documents: once new snippets are provided, it reindexes them and starts to use the new knowledge to answer subsequent queries.
+
+### Watch a Demo Here
+(Available soon)
+
+
+### Key Features
+- **HTTP REST queries:** The system is capable of responding in real-time to HTTP REST queries.
+- **Real-time document indexing pipeline:** This pipeline reads data directly from S3-compatible storage, without the need to query a vector document database.
+- **User session and beta testing handling:** The query building process can be extended to handle user sessions and beta testing for new models.
+- **Code reusability for offline evaluation:** The same code can be used for static evaluation of the system.
+
+## Getting Started
+
+### Installation
+
+Clone the repository and `cd` to it. Create a new environment and install the required packages:
+
+```bash
+python -m venv pw-env && source pw-env/bin/activate
+pip install --upgrade --extra-index-url https://packages.pathway.com/966431ef6ba -r requirements.txt
+```
+### Usage
+
+- Create an .env file and add the following environment variables:
+```bash
+PATHWAY_REST_CONNECTOR_HOST=127.0.0.1
+PATHWAY_REST_CONNECTOR_PORT=8080
+OPENAI_API_TOKEN=<Your Token>
+PATHWAY_CACHE_DIR=/tmp/cache
+```
+
+- Run the script using the command: 
+```bash 
+cd llm-app/
+python main.py --mode contextful
+```
+You can also run the app without the need for external APIs by using `local` mode.
+
+- Send REST queries (in a separate terminal window):
+```bash
+curl --data '{"user": "user", "query": "How to connect to Kafka in Pathway?"}' http://localhost:8080/ | jq
+
+curl --data '{"user": "user", "query": "How to use LLMs in Pathway?"}' http://localhost:8080/ | jq
+```
+
+- Test reactivity by adding a new file:
+```bash
+cp ./data/documents_extra.jsonl ./data/pathway-docs/
+curl --data '{"user": "user", "query": "How to use LLMs in Pathway?"}' http://localhost:8080/ | jq
+```
+
+## Further Reading
+Read more about the implementation details and how to extend this application in our [blog series](https://pathway.com/blog/?tag=tutorial).

data/.gitignore

@@ -0,0 +1,4 @@
+llm/
+.vscode/
+__pycache__/
+**/.env

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Pathway
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,80 @@
+
+<div align="center">
+  <img src="https://pathway.com/logo-light.svg" /><br /><br />
+</div>
+<p align="center">
+    <a href="https://github.com/pathwaycom/llm-app-pathway/blob/main/LICENSE">
+        <img src="https://img.shields.io/github/license/pathwaycom/llm-app-pathway?style=plastic" alt="Contributors"/></a>
+    <a href="https://github.com/pathwaycom/llm-app-pathway/graphs/contributors">
+        <img src="https://img.shields.io/github/contributors/pathwaycom/llm-app-pathway?style=plastic" alt="Contributors"/></a>
+    <a href="https://github.com/pathwaycom/llm-app-pathway/actions/workflows/install_package.yml">
+        <img src="https://img.shields.io/github/actions/workflow/status/pathwaycom/llm-app-pathway/install_package.yml?style=plastic" alt="Build" /></a> 
+        <img src="https://img.shields.io/badge/OS-Linux-green" alt="Linux"/>
+        <img src="https://img.shields.io/badge/OS-macOS-green" alt="macOS"/>
+      <br>
+    <a href="https://discord.gg/pathway">
+        <img src="https://img.shields.io/discord/1042405378304004156?logo=discord"
+            alt="chat on Discord"></a>
+    <a href="https://twitter.com/intent/follow?screen_name=pathway_com">
+        <img src="https://img.shields.io/twitter/follow/pathway_com?style=social&logo=twitter"
+            alt="follow on Twitter"></a>
+</p>
+
+This repository contains the code for our Pathway-based LLM App, which is designed to provide real-time responses to queries about the Pathway documentation.
+
+## Project Overview
+
+The app reads a corpus of documents stored in S3 or locally, preprocesses them, and builds a Pathway vector index. It then listens to user queries coming as HTTP REST requests. Each query uses the index to retrieve relevant documentation snippets and uses the OpenAI API/ Hugging Face to provide a response in natural language. The bot is reactive to changes in the corpus of documents: once new snippets are provided, it reindexes them and starts to use the new knowledge to answer subsequent queries.
+
+### Watch a Demo Here
+(Available soon)
+
+
+### Key Features
+- **HTTP REST queries:** The system is capable of responding in real-time to HTTP REST queries.
+- **Real-time document indexing pipeline:** This pipeline reads data directly from S3-compatible storage, without the need to query a vector document database.
+- **User session and beta testing handling:** The query building process can be extended to handle user sessions and beta testing for new models.
+- **Code reusability for offline evaluation:** The same code can be used for static evaluation of the system.
+
+## Getting Started
+
+### Installation
+
+Clone the repository and `cd` to it. Create a new environment and install the required packages:
+
+```bash
+python -m venv pw-env && source pw-env/bin/activate
+pip install --upgrade --extra-index-url https://packages.pathway.com/966431ef6ba -r requirements.txt
+```
+### Usage
+
+- Create an .env file and add the following environment variables:
+```bash
+PATHWAY_REST_CONNECTOR_HOST=127.0.0.1
+PATHWAY_REST_CONNECTOR_PORT=8080
+OPENAI_API_TOKEN=<Your Token>
+PATHWAY_CACHE_DIR=/tmp/cache
+```
+
+- Run the script using the command: 
+```bash 
+cd llm-app/
+python main.py --mode contextful
+```
+You can also run the app without the need for external APIs by using `local` mode.
+
+- Send REST queries (in a separate terminal window):
+```bash
+curl --data '{"user": "user", "query": "How to connect to Kafka in Pathway?"}' http://localhost:8080/ | jq
+
+curl --data '{"user": "user", "query": "How to use LLMs in Pathway?"}' http://localhost:8080/ | jq
+```
+
+- Test reactivity by adding a new file:
+```bash
+cp ./data/documents_extra.jsonl ./data/pathway-docs/
+curl --data '{"user": "user", "query": "How to use LLMs in Pathway?"}' http://localhost:8080/ | jq
+```
+
+## Further Reading
+Read more about the implementation details and how to extend this application in our [blog series](https://pathway.com/blog/?tag=tutorial).

data/documents_extra.jsonl

@@ -0,0 +1 @@
+{"doc": "Using Large Language Models in Pathway is simple: just call the functions from `pathway.stdlib.ml.nlp`!"}

data/llm-app/main.py

@@ -0,0 +1,19 @@
+import argparse
+from dotenv import load_dotenv
+import importlib
+
+
+load_dotenv()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="LLM App Pathway")
+
+    parser.add_argument("--mode", type=str,
+                        choices=['contextful', 'contextless', 'local'],
+                        default='contextful', help="Which pathway logic to run (default: %(default)s)")
+
+    args = parser.parse_args()
+    scenario_module = importlib.import_module(f'pathway_pipelines.{args.mode}.app')
+
+    scenario_module.run()

data/llm-app/model_wrappers/__init__.py

@@ -0,0 +1,3 @@
+from model_wrappers.openai_wrapper.api_models import OpenAIEmbeddingModel, OpenAIChatGPTModel
+from model_wrappers.huggingface_wrapper.api_models import HFApiFeatureExtractionTask, HFApiTextGenerationTask
+from model_wrappers.huggingface_wrapper.pipelines import HFFeatureExtractionTask, HFTextGenerationTask

data/llm-app/model_wrappers/api_clients/clients.py

@@ -0,0 +1,42 @@
+import logging
+from abc import ABC, abstractmethod
+import openai
+import requests
+
+logfun = logging.debug
+
+
+class APIClient(ABC):
+    @abstractmethod
+    def make_request(self, **kwargs):
+        pass
+
+
+class OpenAIClient(APIClient):
+    def __init__(self, api_key: str):
+        openai.api_key = api_key
+
+
+class OpenAIChatCompletionClient(OpenAIClient):
+    def make_request(self, **kwargs):
+        logfun("Calling OpenAI chat completion service %s", str(kwargs)[:100])
+        return openai.ChatCompletion.create(**kwargs)
+
+
+class OpenAIEmbeddingClient(OpenAIClient):
+    def make_request(self, **kwargs):
+        logfun("Calling OpenAI embedding service %s", str(kwargs)[:100])
+        return openai.Embedding.create(**kwargs)
+
+
+class HuggingFaceClient(APIClient):
+    def __init__(self, api_key: str) -> None:
+        self.headers = {"Authorization": f"Bearer {api_key}"}
+        self.api_url_prefix = "https://api-inference.huggingface.co/models"
+
+    def make_request(self, **kwargs):
+        logfun("Calling HuggingFace %s", str(kwargs)[:100])
+        endpoint = kwargs.pop("model")
+        url = f"{self.api_url_prefix}/{endpoint}"
+        response = requests.post(url, headers=self.headers, json=kwargs)
+        return response.json()

data/llm-app/model_wrappers/base.py

@@ -0,0 +1,56 @@
+import pathway as pw
+from pathway.internals import expression as expr
+import functools
+import os
+from abc import ABC, abstractmethod
+import diskcache
+from model_wrappers.api_clients.clients import APIClient
+
+
+class _Cache:
+    """A simple cache"""
+
+    def __init__(self) -> None:
+        if cache_dir := os.environ.get("PATHWAY_CACHE_DIR"):
+            self.cache = diskcache.Cache(cache_dir)
+        else:
+            self.cache = {}
+
+    def __call__(self, fun):
+        base_name = f"{fun.__module__}_{fun.__qualname__}"
+
+        @functools.wraps(fun)
+        def wrapper(*args, **kwargs):
+            key = f"{base_name}({(args, kwargs)})"
+            if key not in self.cache:
+                self.cache[key] = fun(*args, **kwargs)
+            return self.cache[key]
+
+        return wrapper
+
+
+class BaseModel(ABC):
+    def __init__(self, **kwargs):
+        self.config = kwargs
+        self.cache = _Cache()
+
+    def __call__(self, text: str, **kwargs):
+        raise NotImplementedError()
+
+    def apply(
+        self,
+        text: expr.ColumnExpressionOrConst,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+
+        return pw.apply_async(self, text=text, **kwargs)
+
+
+class APIModel(BaseModel):
+    def __init__(self, api_key: str, **kwargs):
+        super().__init__(**kwargs)
+        self.api_client = self.get_client(api_key)
+
+    @abstractmethod
+    def get_client(self, api_key: str) -> APIClient:
+        pass

data/llm-app/model_wrappers/huggingface_wrapper/api_models.py

@@ -0,0 +1,95 @@
+from api_clients.clients import HuggingFaceClient
+from base import APIModel
+
+
+class HuggingFaceAPIModel(APIModel):
+    def get_client(self, api_key: str) -> HuggingFaceClient:
+        return HuggingFaceClient(api_key=api_key)
+
+    def call_api(self, **kwargs):
+        """
+        Makes a request to the Hugging Face API and returns the result.
+
+        The method accepts arguments as keyword arguments (**kwargs).
+        The expected arguments are 'model' and others that depend on the specific task.
+        Please check [HuggingFace Inference API](https://huggingface.co/docs/api-inference/detailed_parameters)
+        'model' is a string representing the pre-trained model to use for the call.
+
+        Examples:
+
+        1) Question-Answering Task usage:
+
+        model = HuggingFaceModel(api_key=token)
+        result = model.call(
+            inputs=dict(
+                context="Pathway is a realtime stream data processing framework. It has a python api",
+                question="Does Pathway have a python API ?"
+            ),
+            model='deepset/roberta-base-squad2'
+        )
+        The expected output for the example above is:
+        {
+            'score': 0.41046786308288574,
+            'start': 56,
+            'end': 75,
+            'answer': 'It has a python api'
+        }
+
+        2) Sentiment analysis model usage:
+
+        result = model.call(
+            inputs="What a performance that was!",
+            model='distilbert-base-uncased-finetuned-sst-2-english'
+        )
+        The expected output for the example above is:
+        [
+            [
+                {'label': 'POSITIVE', 'score': 0.9988183379173279},
+                {'label': 'NEGATIVE', 'score': 0.0011816363548859954}
+            ]
+        ]
+
+        Args:
+            **kwargs: The arguments for the model call. 'inputs' and 'model' keys are expected.
+
+        Returns:
+            The response from the Hugging Face API, the format depends on the model being used.
+        """
+        return self.api_client.make_request(**kwargs)
+
+
+class HFApiFeatureExtractionTask(HuggingFaceAPIModel):
+    def __call__(
+        self, text: str, locator="sentence-transformers/all-MiniLM-L6-v2", **kwargs
+    ):
+        response = self.call_api(inputs=[text], model=locator, **kwargs)
+        return response
+
+
+class HFApiTextGenerationTask(HuggingFaceAPIModel):
+    """
+    A class that represents a text generation task using the Hugging Face API.
+
+    It inherits from the HuggingFaceModel class and overrides the call method to
+    specifically work with text generation models.
+
+    This class allows users to simply pass a text string and get a generated
+    text in return.
+
+    Args:
+        api_key (str): The API key to access the Hugging Face API.
+
+    Example:
+
+    # >>> model = HFTextGenerationTask(api_key=token)
+    # >>> text = "Once upon a time"
+    # >>> generated_text = model(text, locator="gpt2")
+    # >>> print(generated_text)
+    'Once upon a time in a land far away...'
+
+    """
+
+    def __call__(self, text: str, locator="gpt2", **kwargs):
+        response = self.call_api(inputs=text, model=locator, **kwargs)
+        return response[0]["generated_text"]
+

data/llm-app/model_wrappers/huggingface_wrapper/pipelines.py

@@ -0,0 +1,49 @@
+from transformers import pipeline
+from model_wrappers.base import BaseModel
+
+
+class HFPipelineTask(BaseModel):
+    def __init__(self, model_name, device="cpu", **kwargs):
+        super().__init__(**kwargs)
+        self.pipeline = pipeline(
+            model=model_name, device=device
+        )
+        self.tokenizer = self.pipeline.tokenizer
+
+    def crop_to_max_length(self, input_string, max_length=500):
+        tokens = self.tokenizer.tokenize(input_string)
+        if len(tokens) > max_length:
+            tokens = tokens[:max_length]
+        return self.tokenizer.convert_tokens_to_string(tokens)
+
+
+class HFFeatureExtractionTask(HFPipelineTask):
+    def __init__(self, max_length=500, **kwargs):
+        super().__init__(**kwargs)
+        self.max_length = max_length
+
+    def __call__(self, text, **kwargs):
+
+        text = self.crop_to_max_length(text, max_length=self.max_length)
+        # This will return a list of lists (one list for each word in the text)
+        embedding = self.pipeline(text, **kwargs)[0]
+
+        # For simplicity, we'll just average all word vectors to get a sentence embedding
+        avg_embedding = [sum(col) / len(col) for col in zip(*embedding)]
+
+        return avg_embedding
+
+
+class HFTextGenerationTask(HFPipelineTask):
+    def __init__(self, max_prompt_length=500, max_new_tokens=500, **kwargs):
+        super().__init__(**kwargs)
+        self.max_prompt_length = max_prompt_length
+        self.max_new_tokens = max_new_tokens
+
+    def __call__(self, text, **kwargs):
+        text = self.crop_to_max_length(text, self.max_prompt_length)
+
+        max_new_tokens = kwargs.pop('max_new_tokens', self.max_new_tokens)
+
+        output = self.pipeline(text, max_new_tokens=max_new_tokens, **kwargs)
+        return output[0]["generated_text"]

data/llm-app/model_wrappers/openai_wrapper/api_models.py

@@ -0,0 +1,151 @@
+from pathway.internals import expression as expr
+from model_wrappers.api_clients.clients import (
+    OpenAIChatCompletionClient,
+    OpenAIClient,
+    OpenAIEmbeddingClient,
+)
+from model_wrappers.base import APIModel
+
+
+class MessagePreparer:
+    @staticmethod
+    def prepare_chat_messages(prompt: str):
+        return [
+            dict(role="system", content="You are a helpful assistant"),
+            dict(role="user", content=prompt),
+        ]
+
+
+class OpenAIChatGPTModel(APIModel):
+    def get_client(self, openai_key: str) -> OpenAIClient:
+        return OpenAIChatCompletionClient(openai_key)
+
+    def __call__(self, text: str, locator="gpt-3.5-turbo", **kwargs) -> str:
+        """
+
+        Example
+
+        # >>> model = OpenAIChatGPTModel(api_key = api_key)
+        # >>> model(
+        # ...     locator='gpt-4-0613',
+        # ...     text="Tell me a joke about jokes",
+        # ...     temperature=1.1
+        # ... )
+        """
+        messages = MessagePreparer.prepare_chat_messages(text)
+        response = self.api_client.make_request(
+            messages=messages, model=locator, **kwargs
+        )
+        return response.choices[0].message.content
+
+    def apply(
+        self,
+        *args,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+        """
+        Applies the specified model in `locator` from OpenAIChatGPT API to the provided text.
+        Parameters
+        ----------
+        text : Union[pw.ColumnExpression, str]
+            The input text on which the model will be applied. It can be a column expression or a string.
+        locator : Union[pw.ColumnExpression, str, None]
+            The model locator to use for applying the model.
+            If provided, it should be a column expression or a string.
+            Otherwise, the default chat completion model `gpt-3.5-turbo` is applied.
+            Please check out https://platform.openai.com/docs/models/model-endpoint-compatibility
+            to see the available models.
+        **kwargs : dict
+            Additional keyword arguments that will be used for the model application.
+            These could include settings such as `temperature`, `max_tokens`, etc.
+            Check https://platform.openai.com/docs/api-reference/chat/create for the official API Reference
+        Returns
+        -------
+        pw.ColumnExpression
+            The result of the model application as a column expression or str.
+            Please note that the output is `chat_completion.choices[0].message.content`
+            where `chat_completion` is the api response.
+        Example:
+        # >>> model = OpenAIChatGPTModel(api_key = api_key)
+        # >>>
+        # >>> table = pw.debug.table_from_pandas(
+        # ...     pd.DataFrame.from_records([
+        # ...         {"text": "How to use pathway to process a kafka stream ?"},
+        # ...         {"text": "How to apply a function to a pathway table ?"}
+        # ...     ])
+        # ... )
+        # >>> table += table.select(
+        # ...     response = model.apply(
+        # ...         pw.this.text,
+        # ...         locator='gpt-4',
+        # ...         temperature=1.5,
+        # ...         max_tokens=1000
+        # ...     )
+        # ... )
+        """
+        return super().apply(*args, **kwargs)
+
+
+class OpenAIEmbeddingModel(APIModel):
+    def get_client(self, openai_key: str) -> OpenAIClient:
+        return OpenAIEmbeddingClient(openai_key)
+
+    def __call__(self, text: str, locator="text-embedding-ada-002", **kwargs):
+        """
+        Example:
+
+        # >>> embedder = OpenAIEmbeddingModel(api_key)
+        # >>>
+        # >>> embedder(
+        # ...    text='Some random text'
+        # ...    locator='text-embedding-ada-002'
+        # ... )
+        """
+
+        response = self.api_client.make_request(input=[text], model=locator, **kwargs)
+        return response["data"][0]["embedding"]
+
+    def apply(
+        self,
+        *args,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+        """
+        Applies the specified model in `locator` from OpenAIEmbeddingModel API to the provided text.
+        Parameters
+        ----------
+        text : Union[pw.ColumnExpression, str]
+            The input text on which the model will be applied. It can be a column expression or a constant value.
+        locator : Union[pw.ColumnExpression, str, None]
+            The model locator to use for applying the model.
+            If provided, it should be a column expression or a constant value.
+            Otherwise, the default chat completion model `gpt-3.5-turbo` is applied.
+            Please check out https://platform.openai.com/docs/models/model-endpoint-compatibility
+            to see the available models.
+        **kwargs : dict
+            Additional keyword arguments that will be used for the model application.
+            These could include settings such as `temperature`, `max_tokens`, etc.
+            You can check https://platform.openai.com/docs/api-reference/embeddings/create
+            for the official API Reference.
+        Returns
+        -------
+        pw.ColumnExpression
+            The result of the model application as a column expression or constant of type list.
+            Please note that the output is `results["data"][0]["embedding"]`
+        Example:
+        # >>> embedder = OpenAIEmbeddingModel(api_key)
+        # >>>
+        # >>> table = pw.debug.table_from_pandas(
+        # ...     pd.DataFrame.from_records([
+        # ...         {"text": "How to use pathway to process a kafka stream ?"},
+        # ...         {"text": "How to apply a function to a pathway table ?"}
+        # ...     ])
+        # ... )
+        # >>> table += table.select(
+        # ...     embedding = embedder.apply(
+        # ...         pw.this.text,
+        # ...         locator='text-embedding-ada-002'
+        # ...     )
+        # ... )
+        """
+        return super().apply(*args, **kwargs)

data/llm-app/pathway_pipelines/contextful/app.py

@@ -0,0 +1,105 @@
+"""
+Microservice for  a context-aware ChatGPT assistant.
+
+The following program reads in a collection of documents,
+embeds each document using the OpenAI document embedding model,
+then builds an index for fast retrieval of documents relevant to a question,
+effectively replacing a vector database.
+
+The program then starts a REST API endpoint serving queries about programming in Pathway.
+
+Each query text is first turned into a vector using OpenAI embedding service,
+then relevant documentation pages are found using a Nearest Neighbor index computed
+for documents in the corpus. A prompt is build from the relevant documentations pages
+and sent to the OpenAI GPT-4 chat service for processing.
+
+Usage:
+In llm-app/ run:
+python main.py --mode contextful
+
+To call the REST API:
+curl --data '{"user": "user", "query": "How to connect to Kafka in Pathway?"}' http://localhost:8080/ | jq
+"""
+import os
+import pathway as pw
+from pathway.stdlib.ml.index import KNNIndex
+from model_wrappers import OpenAIEmbeddingModel, OpenAIChatGPTModel
+
+
+class DocumentInputSchema(pw.Schema):
+    doc: str
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+
+API_KEY = os.environ.get("OPENAI_API_TOKEN")
+EMBEDDER_LOCATOR = "text-embedding-ada-002"
+EMBEDDING_DIMENSION = 1536
+
+MODEL_LOCATOR = "gpt-4"
+TEMPERATURE = 0.0
+MAX_TOKENS = 60
+
+
+def run():
+    embedder = OpenAIEmbeddingModel(api_key=API_KEY)
+
+    documents = pw.io.jsonlines.read(
+        "../data/pathway-docs/",
+        schema=DocumentInputSchema,
+        mode="streaming",
+        autocommit_duration_ms=50,
+    )
+
+    enriched_documents = documents + documents.select(
+        data=embedder.apply(text=pw.this.doc, locator=EMBEDDER_LOCATOR)
+    )
+
+    index = KNNIndex(enriched_documents, d=EMBEDDING_DIMENSION)
+
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    query += query.select(
+        data=embedder.apply(text=pw.this.query, locator=EMBEDDER_LOCATOR),
+    )
+
+    query_context = index.query(query, k=3).select(
+        pw.this.query, documents_list=pw.this.result
+    )
+
+    @pw.udf
+    def build_prompt(documents, query):
+        docs_str = "\n".join(documents)
+        prompt = f"Given the following documents : \n {docs_str} \nanswer this query: {query}"
+        return prompt
+
+    prompt = query_context.select(
+        prompt=build_prompt(pw.this.documents_list, pw.this.query)
+    )
+
+    model = OpenAIChatGPTModel(api_key=API_KEY)
+
+    responses = prompt.select(
+        query_id=pw.this.id,
+        result=model.apply(
+            pw.this.prompt,
+            locator=MODEL_LOCATOR,
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS
+        ),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

data/llm-app/pathway_pipelines/contextless/app.py

@@ -0,0 +1,41 @@
+import os
+import pathway as pw
+from model_wrappers import OpenAIChatGPTModel
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+API_KEY = os.environ.get("OPENAI_API_TOKEN")
+MODEL_LOCATOR = "gpt-4"
+TEMPERATURE = 0.0
+MAX_TOKENS = 50
+
+
+def run():
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    model = OpenAIChatGPTModel(api_key=API_KEY)
+
+    responses = query.select(
+        query_id=pw.this.id,
+        result=model.apply(
+            pw.this.query,
+            locator=MODEL_LOCATOR,
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS
+        ),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

data/llm-app/pathway_pipelines/local/app.py

@@ -0,0 +1,76 @@
+import os
+import pathway as pw
+from pathway.stdlib.ml.index import KNNIndex
+from model_wrappers import (
+    HFFeatureExtractionTask,
+    HFTextGenerationTask,
+)
+
+
+class DocumentInputSchema(pw.Schema):
+    doc: str
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+
+EMBEDDER_LOCATOR = "intfloat/e5-large-v2"
+EMBEDDING_DIMENSION = 1024
+MODEL_LOCATOR = "gpt2"
+
+
+def run():
+    embedder = HFFeatureExtractionTask(model_name=EMBEDDER_LOCATOR)
+
+    documents = pw.io.jsonlines.read(
+        "../data/pathway-docs/", schema=DocumentInputSchema, mode="streaming", autocommit_duration_ms=50
+    )
+
+    enriched_documents = documents + documents.select(
+        data=embedder.apply(text=pw.this.doc)
+    )
+
+    index = KNNIndex(enriched_documents, d=EMBEDDING_DIMENSION)
+
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    query += query.select(
+        data=embedder.apply(text=pw.this.query),
+    )
+
+    query_context = index.query(query, k=1).select(
+        pw.this.query, documents_list=pw.this.result
+    )
+
+    @pw.udf
+    def build_prompt(documents, query):
+        docs_str = "\n".join(documents)
+        prompt = (
+            f"Given the following documents : \n {docs_str} \nanswer this query: {query}"
+        )
+        return prompt
+
+    prompt = query_context.select(
+        prompt=build_prompt(pw.this.documents_list, pw.this.query)
+    )
+
+    model = HFTextGenerationTask(model_name=MODEL_LOCATOR)
+
+    responses = prompt.select(
+        query_id=pw.this.id,
+        result=model.apply(pw.this.prompt, return_full_text=False, max_new_tokens=60),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

data/requirements.txt

@@ -0,0 +1,6 @@
+pathway
+openai
+transformers[torch]
+requests
+diskcache
+python-dotenv

llm-app/main.py

@@ -0,0 +1,19 @@
+import argparse
+from dotenv import load_dotenv
+import importlib
+
+
+load_dotenv()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="LLM App Pathway")
+
+    parser.add_argument("--mode", type=str,
+                        choices=['contextful', 'contextless', 'local'],
+                        default='contextful', help="Which pathway logic to run (default: %(default)s)")
+
+    args = parser.parse_args()
+    scenario_module = importlib.import_module(f'pathway_pipelines.{args.mode}.app')
+
+    scenario_module.run()

llm-app/model_wrappers/__init__.py

@@ -0,0 +1,3 @@
+from model_wrappers.openai_wrapper.api_models import OpenAIEmbeddingModel, OpenAIChatGPTModel
+from model_wrappers.huggingface_wrapper.api_models import HFApiFeatureExtractionTask, HFApiTextGenerationTask
+from model_wrappers.huggingface_wrapper.pipelines import HFFeatureExtractionTask, HFTextGenerationTask

llm-app/model_wrappers/api_clients/clients.py

@@ -0,0 +1,42 @@
+import logging
+from abc import ABC, abstractmethod
+import openai
+import requests
+
+logfun = logging.debug
+
+
+class APIClient(ABC):
+    @abstractmethod
+    def make_request(self, **kwargs):
+        pass
+
+
+class OpenAIClient(APIClient):
+    def __init__(self, api_key: str):
+        openai.api_key = api_key
+
+
+class OpenAIChatCompletionClient(OpenAIClient):
+    def make_request(self, **kwargs):
+        logfun("Calling OpenAI chat completion service %s", str(kwargs)[:100])
+        return openai.ChatCompletion.create(**kwargs)
+
+
+class OpenAIEmbeddingClient(OpenAIClient):
+    def make_request(self, **kwargs):
+        logfun("Calling OpenAI embedding service %s", str(kwargs)[:100])
+        return openai.Embedding.create(**kwargs)
+
+
+class HuggingFaceClient(APIClient):
+    def __init__(self, api_key: str) -> None:
+        self.headers = {"Authorization": f"Bearer {api_key}"}
+        self.api_url_prefix = "https://api-inference.huggingface.co/models"
+
+    def make_request(self, **kwargs):
+        logfun("Calling HuggingFace %s", str(kwargs)[:100])
+        endpoint = kwargs.pop("model")
+        url = f"{self.api_url_prefix}/{endpoint}"
+        response = requests.post(url, headers=self.headers, json=kwargs)
+        return response.json()

llm-app/model_wrappers/base.py

@@ -0,0 +1,56 @@
+import pathway as pw
+from pathway.internals import expression as expr
+import functools
+import os
+from abc import ABC, abstractmethod
+import diskcache
+from model_wrappers.api_clients.clients import APIClient
+
+
+class _Cache:
+    """A simple cache"""
+
+    def __init__(self) -> None:
+        if cache_dir := os.environ.get("PATHWAY_CACHE_DIR"):
+            self.cache = diskcache.Cache(cache_dir)
+        else:
+            self.cache = {}
+
+    def __call__(self, fun):
+        base_name = f"{fun.__module__}_{fun.__qualname__}"
+
+        @functools.wraps(fun)
+        def wrapper(*args, **kwargs):
+            key = f"{base_name}({(args, kwargs)})"
+            if key not in self.cache:
+                self.cache[key] = fun(*args, **kwargs)
+            return self.cache[key]
+
+        return wrapper
+
+
+class BaseModel(ABC):
+    def __init__(self, **kwargs):
+        self.config = kwargs
+        self.cache = _Cache()
+
+    def __call__(self, text: str, **kwargs):
+        raise NotImplementedError()
+
+    def apply(
+        self,
+        text: expr.ColumnExpressionOrConst,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+
+        return pw.apply_async(self, text=text, **kwargs)
+
+
+class APIModel(BaseModel):
+    def __init__(self, api_key: str, **kwargs):
+        super().__init__(**kwargs)
+        self.api_client = self.get_client(api_key)
+
+    @abstractmethod
+    def get_client(self, api_key: str) -> APIClient:
+        pass

llm-app/model_wrappers/huggingface_wrapper/api_models.py

@@ -0,0 +1,95 @@
+from model_wrappers.api_clients.clients import HuggingFaceClient
+from model_wrappers.base import APIModel
+
+
+class HuggingFaceAPIModel(APIModel):
+    def get_client(self, api_key: str) -> HuggingFaceClient:
+        return HuggingFaceClient(api_key=api_key)
+
+    def call_api(self, **kwargs):
+        """
+        Makes a request to the Hugging Face API and returns the result.
+
+        The method accepts arguments as keyword arguments (**kwargs).
+        The expected arguments are 'model' and others that depend on the specific task.
+        Please check [HuggingFace Inference API](https://huggingface.co/docs/api-inference/detailed_parameters)
+        'model' is a string representing the pre-trained model to use for the call.
+
+        Examples:
+
+        1) Question-Answering Task usage:
+
+        model = HuggingFaceModel(api_key=token)
+        result = model.call(
+            inputs=dict(
+                context="Pathway is a realtime stream data processing framework. It has a python api",
+                question="Does Pathway have a python API ?"
+            ),
+            model='deepset/roberta-base-squad2'
+        )
+        The expected output for the example above is:
+        {
+            'score': 0.41046786308288574,
+            'start': 56,
+            'end': 75,
+            'answer': 'It has a python api'
+        }
+
+        2) Sentiment analysis model usage:
+
+        result = model.call(
+            inputs="What a performance that was!",
+            model='distilbert-base-uncased-finetuned-sst-2-english'
+        )
+        The expected output for the example above is:
+        [
+            [
+                {'label': 'POSITIVE', 'score': 0.9988183379173279},
+                {'label': 'NEGATIVE', 'score': 0.0011816363548859954}
+            ]
+        ]
+
+        Args:
+            **kwargs: The arguments for the model call. 'inputs' and 'model' keys are expected.
+
+        Returns:
+            The response from the Hugging Face API, the format depends on the model being used.
+        """
+        return self.api_client.make_request(**kwargs)
+
+
+class HFApiFeatureExtractionTask(HuggingFaceAPIModel):
+    def __call__(
+        self, text: str, locator="sentence-transformers/all-MiniLM-L6-v2", **kwargs
+    ):
+        response = self.call_api(inputs=[text], model=locator, **kwargs)
+        return response
+
+
+class HFApiTextGenerationTask(HuggingFaceAPIModel):
+    """
+    A class that represents a text generation task using the Hugging Face API.
+
+    It inherits from the HuggingFaceModel class and overrides the call method to
+    specifically work with text generation models.
+
+    This class allows users to simply pass a text string and get a generated
+    text in return.
+
+    Args:
+        api_key (str): The API key to access the Hugging Face API.
+
+    Example:
+
+    # >>> model = HFTextGenerationTask(api_key=token)
+    # >>> text = "Once upon a time"
+    # >>> generated_text = model(text, locator="gpt2")
+    # >>> print(generated_text)
+    'Once upon a time in a land far away...'
+
+    """
+
+    def __call__(self, text: str, locator="gpt2", **kwargs):
+        response = self.call_api(inputs=text, model=locator, **kwargs)
+        return response[0]["generated_text"]
+

llm-app/model_wrappers/huggingface_wrapper/pipelines.py

@@ -0,0 +1,49 @@
+from transformers import pipeline
+from model_wrappers.base import BaseModel
+
+
+class HFPipelineTask(BaseModel):
+    def __init__(self, model_name, device="cpu", **kwargs):
+        super().__init__(**kwargs)
+        self.pipeline = pipeline(
+            model=model_name, device=device
+        )
+        self.tokenizer = self.pipeline.tokenizer
+
+    def crop_to_max_length(self, input_string, max_length=500):
+        tokens = self.tokenizer.tokenize(input_string)
+        if len(tokens) > max_length:
+            tokens = tokens[:max_length]
+        return self.tokenizer.convert_tokens_to_string(tokens)
+
+
+class HFFeatureExtractionTask(HFPipelineTask):
+    def __init__(self, max_length=500, **kwargs):
+        super().__init__(**kwargs)
+        self.max_length = max_length
+
+    def __call__(self, text, **kwargs):
+
+        text = self.crop_to_max_length(text, max_length=self.max_length)
+        # This will return a list of lists (one list for each word in the text)
+        embedding = self.pipeline(text, **kwargs)[0]
+
+        # For simplicity, we'll just average all word vectors to get a sentence embedding
+        avg_embedding = [sum(col) / len(col) for col in zip(*embedding)]
+
+        return avg_embedding
+
+
+class HFTextGenerationTask(HFPipelineTask):
+    def __init__(self, max_prompt_length=500, max_new_tokens=500, **kwargs):
+        super().__init__(**kwargs)
+        self.max_prompt_length = max_prompt_length
+        self.max_new_tokens = max_new_tokens
+
+    def __call__(self, text, **kwargs):
+        text = self.crop_to_max_length(text, self.max_prompt_length)
+
+        max_new_tokens = kwargs.pop('max_new_tokens', self.max_new_tokens)
+
+        output = self.pipeline(text, max_new_tokens=max_new_tokens, **kwargs)
+        return output[0]["generated_text"]

llm-app/model_wrappers/openai_wrapper/api_models.py

@@ -0,0 +1,151 @@
+from pathway.internals import expression as expr
+from model_wrappers.api_clients.clients import (
+    OpenAIChatCompletionClient,
+    OpenAIClient,
+    OpenAIEmbeddingClient,
+)
+from model_wrappers.base import APIModel
+
+
+class MessagePreparer:
+    @staticmethod
+    def prepare_chat_messages(prompt: str):
+        return [
+            dict(role="system", content="You are a helpful assistant"),
+            dict(role="user", content=prompt),
+        ]
+
+
+class OpenAIChatGPTModel(APIModel):
+    def get_client(self, openai_key: str) -> OpenAIClient:
+        return OpenAIChatCompletionClient(openai_key)
+
+    def __call__(self, text: str, locator="gpt-3.5-turbo", **kwargs) -> str:
+        """
+
+        Example
+
+        # >>> model = OpenAIChatGPTModel(api_key = api_key)
+        # >>> model(
+        # ...     locator='gpt-4-0613',
+        # ...     text="Tell me a joke about jokes",
+        # ...     temperature=1.1
+        # ... )
+        """
+        messages = MessagePreparer.prepare_chat_messages(text)
+        response = self.api_client.make_request(
+            messages=messages, model=locator, **kwargs
+        )
+        return response.choices[0].message.content
+
+    def apply(
+        self,
+        *args,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+        """
+        Applies the specified model in `locator` from OpenAIChatGPT API to the provided text.
+        Parameters
+        ----------
+        text : Union[pw.ColumnExpression, str]
+            The input text on which the model will be applied. It can be a column expression or a string.
+        locator : Union[pw.ColumnExpression, str, None]
+            The model locator to use for applying the model.
+            If provided, it should be a column expression or a string.
+            Otherwise, the default chat completion model `gpt-3.5-turbo` is applied.
+            Please check out https://platform.openai.com/docs/models/model-endpoint-compatibility
+            to see the available models.
+        **kwargs : dict
+            Additional keyword arguments that will be used for the model application.
+            These could include settings such as `temperature`, `max_tokens`, etc.
+            Check https://platform.openai.com/docs/api-reference/chat/create for the official API Reference
+        Returns
+        -------
+        pw.ColumnExpression
+            The result of the model application as a column expression or str.
+            Please note that the output is `chat_completion.choices[0].message.content`
+            where `chat_completion` is the api response.
+        Example:
+        # >>> model = OpenAIChatGPTModel(api_key = api_key)
+        # >>>
+        # >>> table = pw.debug.table_from_pandas(
+        # ...     pd.DataFrame.from_records([
+        # ...         {"text": "How to use pathway to process a kafka stream ?"},
+        # ...         {"text": "How to apply a function to a pathway table ?"}
+        # ...     ])
+        # ... )
+        # >>> table += table.select(
+        # ...     response = model.apply(
+        # ...         pw.this.text,
+        # ...         locator='gpt-4',
+        # ...         temperature=1.5,
+        # ...         max_tokens=1000
+        # ...     )
+        # ... )
+        """
+        return super().apply(*args, **kwargs)
+
+
+class OpenAIEmbeddingModel(APIModel):
+    def get_client(self, openai_key: str) -> OpenAIClient:
+        return OpenAIEmbeddingClient(openai_key)
+
+    def __call__(self, text: str, locator="text-embedding-ada-002", **kwargs):
+        """
+        Example:
+
+        # >>> embedder = OpenAIEmbeddingModel(api_key)
+        # >>>
+        # >>> embedder(
+        # ...    text='Some random text'
+        # ...    locator='text-embedding-ada-002'
+        # ... )
+        """
+
+        response = self.api_client.make_request(input=[text], model=locator, **kwargs)
+        return response["data"][0]["embedding"]
+
+    def apply(
+        self,
+        *args,
+        **kwargs,
+    ) -> expr.ColumnExpressionOrConst:
+        """
+        Applies the specified model in `locator` from OpenAIEmbeddingModel API to the provided text.
+        Parameters
+        ----------
+        text : Union[pw.ColumnExpression, str]
+            The input text on which the model will be applied. It can be a column expression or a constant value.
+        locator : Union[pw.ColumnExpression, str, None]
+            The model locator to use for applying the model.
+            If provided, it should be a column expression or a constant value.
+            Otherwise, the default chat completion model `gpt-3.5-turbo` is applied.
+            Please check out https://platform.openai.com/docs/models/model-endpoint-compatibility
+            to see the available models.
+        **kwargs : dict
+            Additional keyword arguments that will be used for the model application.
+            These could include settings such as `temperature`, `max_tokens`, etc.
+            You can check https://platform.openai.com/docs/api-reference/embeddings/create
+            for the official API Reference.
+        Returns
+        -------
+        pw.ColumnExpression
+            The result of the model application as a column expression or constant of type list.
+            Please note that the output is `results["data"][0]["embedding"]`
+        Example:
+        # >>> embedder = OpenAIEmbeddingModel(api_key)
+        # >>>
+        # >>> table = pw.debug.table_from_pandas(
+        # ...     pd.DataFrame.from_records([
+        # ...         {"text": "How to use pathway to process a kafka stream ?"},
+        # ...         {"text": "How to apply a function to a pathway table ?"}
+        # ...     ])
+        # ... )
+        # >>> table += table.select(
+        # ...     embedding = embedder.apply(
+        # ...         pw.this.text,
+        # ...         locator='text-embedding-ada-002'
+        # ...     )
+        # ... )
+        """
+        return super().apply(*args, **kwargs)

llm-app/pathway_pipelines/contextful/app.py

@@ -0,0 +1,105 @@
+"""
+Microservice for  a context-aware ChatGPT assistant.
+
+The following program reads in a collection of documents,
+embeds each document using the OpenAI document embedding model,
+then builds an index for fast retrieval of documents relevant to a question,
+effectively replacing a vector database.
+
+The program then starts a REST API endpoint serving queries about programming in Pathway.
+
+Each query text is first turned into a vector using OpenAI embedding service,
+then relevant documentation pages are found using a Nearest Neighbor index computed
+for documents in the corpus. A prompt is build from the relevant documentations pages
+and sent to the OpenAI GPT-4 chat service for processing.
+
+Usage:
+In llm-app/ run:
+python main.py --mode contextful
+
+To call the REST API:
+curl --data '{"user": "user", "query": "How to connect to Kafka in Pathway?"}' http://localhost:8080/ | jq
+"""
+import os
+import pathway as pw
+from pathway.stdlib.ml.index import KNNIndex
+from model_wrappers import OpenAIEmbeddingModel, OpenAIChatGPTModel
+
+
+class DocumentInputSchema(pw.Schema):
+    doc: str
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+
+API_KEY = os.environ.get("OPENAI_API_TOKEN")
+EMBEDDER_LOCATOR = "text-embedding-ada-002"
+EMBEDDING_DIMENSION = 1536
+
+MODEL_LOCATOR = "gpt-4"
+TEMPERATURE = 0.0
+MAX_TOKENS = 60
+
+
+def run():
+    embedder = OpenAIEmbeddingModel(api_key=API_KEY)
+
+    documents = pw.io.jsonlines.read(
+        "../data/pathway-docs/",
+        schema=DocumentInputSchema,
+        mode="streaming",
+        autocommit_duration_ms=50,
+    )
+
+    enriched_documents = documents + documents.select(
+        data=embedder.apply(text=pw.this.doc, locator=EMBEDDER_LOCATOR)
+    )
+
+    index = KNNIndex(enriched_documents, d=EMBEDDING_DIMENSION)
+
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    query += query.select(
+        data=embedder.apply(text=pw.this.query, locator=EMBEDDER_LOCATOR),
+    )
+
+    query_context = index.query(query, k=3).select(
+        pw.this.query, documents_list=pw.this.result
+    )
+
+    @pw.udf
+    def build_prompt(documents, query):
+        docs_str = "\n".join(documents)
+        prompt = f"Given the following documents : \n {docs_str} \nanswer this query: {query}"
+        return prompt
+
+    prompt = query_context.select(
+        prompt=build_prompt(pw.this.documents_list, pw.this.query)
+    )
+
+    model = OpenAIChatGPTModel(api_key=API_KEY)
+
+    responses = prompt.select(
+        query_id=pw.this.id,
+        result=model.apply(
+            pw.this.prompt,
+            locator=MODEL_LOCATOR,
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS
+        ),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

llm-app/pathway_pipelines/contextless/app.py

@@ -0,0 +1,41 @@
+import os
+import pathway as pw
+from model_wrappers import OpenAIChatGPTModel
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+API_KEY = os.environ.get("OPENAI_API_TOKEN")
+MODEL_LOCATOR = "gpt-4"
+TEMPERATURE = 0.0
+MAX_TOKENS = 50
+
+
+def run():
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    model = OpenAIChatGPTModel(api_key=API_KEY)
+
+    responses = query.select(
+        query_id=pw.this.id,
+        result=model.apply(
+            pw.this.query,
+            locator=MODEL_LOCATOR,
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS
+        ),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

llm-app/pathway_pipelines/local/app.py

@@ -0,0 +1,76 @@
+import os
+import pathway as pw
+from pathway.stdlib.ml.index import KNNIndex
+from model_wrappers import (
+    HFFeatureExtractionTask,
+    HFTextGenerationTask,
+)
+
+
+class DocumentInputSchema(pw.Schema):
+    doc: str
+
+
+class QueryInputSchema(pw.Schema):
+    query: str
+    user: str
+
+
+HTTP_HOST = os.environ.get("PATHWAY_REST_CONNECTOR_HOST", "127.0.0.1")
+HTTP_PORT = os.environ.get("PATHWAY_REST_CONNECTOR_PORT", "8080")
+
+EMBEDDER_LOCATOR = "intfloat/e5-large-v2"
+EMBEDDING_DIMENSION = 1024
+MODEL_LOCATOR = "gpt2"
+
+
+def run():
+    embedder = HFFeatureExtractionTask(model_name=EMBEDDER_LOCATOR)
+
+    documents = pw.io.jsonlines.read(
+        "../data/pathway-docs/", schema=DocumentInputSchema, mode="streaming", autocommit_duration_ms=50
+    )
+
+    enriched_documents = documents + documents.select(
+        data=embedder.apply(text=pw.this.doc)
+    )
+
+    index = KNNIndex(enriched_documents, d=EMBEDDING_DIMENSION)
+
+    query, response_writer = pw.io.http.rest_connector(
+        host=HTTP_HOST,
+        port=int(HTTP_PORT),
+        schema=QueryInputSchema,
+        autocommit_duration_ms=50,
+    )
+
+    query += query.select(
+        data=embedder.apply(text=pw.this.query),
+    )
+
+    query_context = index.query(query, k=1).select(
+        pw.this.query, documents_list=pw.this.result
+    )
+
+    @pw.udf
+    def build_prompt(documents, query):
+        docs_str = "\n".join(documents)
+        prompt = (
+            f"Given the following documents : \n {docs_str} \nanswer this query: {query}"
+        )
+        return prompt
+
+    prompt = query_context.select(
+        prompt=build_prompt(pw.this.documents_list, pw.this.query)
+    )
+
+    model = HFTextGenerationTask(model_name=MODEL_LOCATOR)
+
+    responses = prompt.select(
+        query_id=pw.this.id,
+        result=model.apply(pw.this.prompt, return_full_text=False, max_new_tokens=60),
+    )
+
+    response_writer(responses)
+
+    pw.run_all(debug=True)

requirements.txt

@@ -0,0 +1,6 @@
+pathway
+openai
+transformers[torch]
+requests
+diskcache
+python-dotenv