Merge pull request #127 from mindflowai/stream

Stream responses from GPT and lazy import CLI commands to improve performance.

Author: steegecs

README.md

@@ -20,12 +20,12 @@ The [ChatGPT](https://openai.com/blog/chatgpt)-powered swiss army knife for the
 ## Getting Started
 
 Pre-requisite: 
-- You'll need to create an [OpenAI](https://openai.com/blog/openai-api) account or request early access from [Anthropic](https://www.anthropic.com/earlyaccess).
+- You'll need to create an [OpenAI](https://openai.com/blog/openai-api) account.
 - Also, create a [Pinecone](https://www.pinecone.io/start) account to use their vector database.
 
 1. Run `pip install mindflow`, or you can clone this repo and run `pip install -e path/to/mindflow`.
 2. Run `mf login`:
-    - Register with OpenAI or Anthropic to use their models. You can find your OpenAI API key [here](https://platform.openai.com/account/api-keys).
+    - Register with OpenAI to use their models. You can find your OpenAI API key [here](https://platform.openai.com/account/api-keys).
     - Register with Pinecone to use their vector database. You can find your Pinecone API key and Environment [here](https://www.pinecone.io/start).
 3. Now, you're ready to start using MindFlow!
 
@@ -42,14 +42,14 @@ There are multiple levels to using mindflow's chat feature.
 - `mf chat "explain what a programming language is"`
     - Interact with chatGPT directly just like on the chatGPT website. We also have chat persistence, so it will remember the previous chat messages.
 2. With File Context
-- `mf chat "please summarize what this code does" path/to/code.py`
+- `mf chat path/to/code.py "please summarize what this code does"`
     - You can provide single or multi-file context to chatGPT by passing in any number of files as a separate argument in the `mf chat` call. For sufficiently small files (see: [chatGPT token limits](https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them)), this will work and also maintain chat history.
 3. With Directory Context
-- `mf chat "what are these submodules responsible for? path/to/submodule1/ path/to/submodule2/`
+- `mf chat path/to/submodule1/ path/to/submodule2/ "what are these submodules responsible for?"`
     - Providing directories will actually run an indexer over your code subdirectories/files recursively. So it may take a while to fully index everything -- don't worry; we'll warn you if the cost becomes a concern! Right now the warning triggers if the index job costs >$0.50USD.
 4. Custom pre-indexed context
 - `mf index path/to/subdir/file1.txt path/to/file2.txt`
-- `mf chat -s "How do all of my classes relate to one another?" ./`
+- `mf chat -s ./ "How do all of my classes relate to one another?"`
     - If you pre-index your repository, you can narrow the scope for the context provided to the chat. Passing `-s` will skip the auto-indexing, and instead will defer to the currently existing index. This index is generated in the first step `mf index` where only those files/subdirs will be included.
     - This can save you time and money if your repository is significantly large.
 
@@ -118,7 +118,7 @@ Make some changes to your branch and stage, and then commit them. Then, run `mf
 ![Screenshot 2023-03-11 at 8 42 11 PM](https://user-images.githubusercontent.com/26421036/224524839-45093b5d-b4d9-4dc4-a129-867d819a2136.png)
 
 ## How does it work?
-This tool allows you to build an index of text documents and search through them using GPT-based embeddings. The tool takes document paths as input, extracts the text, splits the documents into chunks, summarizes them, and builds a summarization tree. The tool then uses this tree to generate embeddings of the indexed documents and your query and selects the top text chunks based on the cosine similarity between these embeddings. The generated index can be saved to a JSON file for later reuse, making subsequent searches faster and cheaper.
+MindFlow uses state-of-the-art methods for high-throughput segmentation, processing, storage, and retrieval of documents using a recursive hierarchical summarization and embedding technique to store embedding vectors for document chunks and then achieve fast, and high-quality responses to questions and tasks by appending similar document chunks based on the hierarchically embedded text and using them as context for you query. Additionally, chat history will persist if it can fit in the context for queries over indexed documents or for regular chat.
 
 ## What's next for MindFlow
 In the future, MindFlow plans on becoming an even more integral part of the modern developer's toolkit. We plan on adding the ability to ditch traditional documentation and instead integrate directly with your private documents and communication channels, allowing for a more seamless and intuitive experience. With MindFlow, you can have a true "stream of consciousness" with your code, documentation, and communication channels, making it easier than ever to stay on top of your projects and collaborate with your team. We are excited to continue pushing the boundaries of what's possible with language models and revolutionizing how developers work.

mindflow/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.3"
+__version__ = "0.5.4"

mindflow/cli/commands/chat.py

@@ -1,21 +1,10 @@
-import os
 import click
-import asyncio
-
 from typing import Tuple
-from result import Result
-
-from mindflow.core.commands.chat import run_chat
-from mindflow.core.commands.index import run_index
-from mindflow.core.commands.query import run_query
-from mindflow.core.settings import Settings
-from mindflow.core.types.model import ModelApiCallError
-from mindflow.core.types.store_traits.json import save_json_store
-from mindflow.core.types.conversation import Conversation
-from mindflow.core.types.definitions.conversation import ConversationID
 
 
 def parse_chat_prompt_and_paths_from_args(prompt_args: Tuple[str]):
+    import os
+
     prompt = " ".join(prompt_args)  # include files/directories in prompt
     paths = []
 
@@ -32,6 +21,34 @@ def parse_chat_prompt_and_paths_from_args(prompt_args: Tuple[str]):
 @click.option("-s", "--skip-index", type=bool, default=False, is_flag=True)
 @click.argument("prompt_args", nargs=-1, type=str, required=True)
 def chat(prompt_args: Tuple[str], skip_index: bool):
+    import click
+    import asyncio
+
+    from typing import List
+    from result import Ok
+
+    from mindflow.core.commands.chat import run_chat
+    from mindflow.core.commands.index import run_index
+    from mindflow.core.commands.query import run_query
+    from mindflow.core.settings import Settings
+    from mindflow.core.types.store_traits.json import save_json_store
+
+    async def stream_chat(settings: Settings, prompt: str):
+        print("\nGPT:")
+        async for char_stream_chunk in run_chat(settings, [], prompt):
+            if isinstance(char_stream_chunk, Ok):
+                click.echo(char_stream_chunk.value, nl=False)
+            else:
+                click.echo(char_stream_chunk.value)
+
+    async def stream_query(settings: Settings, file_paths: List[str], prompt: str):
+        print("\nGPT:")
+        async for char_stream_chunk in run_query(settings, file_paths, prompt):
+            if isinstance(char_stream_chunk, Ok):
+                click.echo(char_stream_chunk.value, nl=False)
+            else:
+                click.echo(char_stream_chunk.value)
+
     prompt, paths = parse_chat_prompt_and_paths_from_args(prompt_args)
     settings = Settings()
     if paths:
@@ -46,18 +63,13 @@ def chat(prompt_args: Tuple[str], skip_index: bool):
 
             asyncio.run(run_index(settings, paths))
 
-        run_query_result: Result[str, ModelApiCallError] = asyncio.run(
-            run_query(settings, paths, prompt)
-        )
-        click.echo(run_query_result.value)
+        asyncio.run(stream_query(settings, paths, prompt))
 
         save_json_store()
         return
 
-    run_chat_result: Result[str, ModelApiCallError] = asyncio.run(
-        run_chat(settings, [], prompt)
-    )
-    click.echo(run_chat_result.value)
+    asyncio.run(stream_chat(settings, prompt))
+
     save_json_store()
 
 
@@ -68,6 +80,9 @@ def history():
 
 @history.command(help="View chat history stats.")
 def stats():
+    from mindflow.core.types.conversation import Conversation
+    from mindflow.core.types.definitions.conversation import ConversationID
+
     if (conversation := Conversation.load(ConversationID.CHAT_0.value)) is None:
         print("No conversation history found.")
         return
@@ -78,6 +93,9 @@ def stats():
 
 @history.command(help="Clear the chat history.")
 def clear():
+    from mindflow.core.types.conversation import Conversation
+    from mindflow.core.types.definitions.conversation import ConversationID
+
     if (conversation := Conversation.load(ConversationID.CHAT_0.value)) is None:
         print("No conversation history found.")
         return

mindflow/cli/commands/config.py

@@ -1,24 +1,115 @@
-import sys
 import click
 from typing import List
 
-from mindflow.core.types.store_traits.json import save_json_store
-from mindflow.core.types.mindflow_model import (
-    MindFlowModel,
-    MindFlowModelConfig,
-    MindFlowModelID,
-)
-
-from mindflow.core.types.definitions.model import (
-    ModelID,
-)
-from mindflow.core.types.model import Model
-
 
 @click.command(
     help="Configure MindFlow. For example, you can configure the model to use."
 )
 def config():
+    from mindflow.core.types.store_traits.json import save_json_store
+    from mindflow.core.types.mindflow_model import (
+        MindFlowModel,
+        MindFlowModelConfig,
+        MindFlowModelID,
+    )
+
+    from mindflow.core.types.definitions.model import (
+        ModelID,
+    )
+    from mindflow.core.types.model import Model
+
+    def configure_model():
+        mindflow_model_ids = [
+            MindFlowModelID.QUERY.value,
+            MindFlowModelID.INDEX.value,
+            MindFlowModelID.EMBEDDING.value,
+        ]
+        mindflow_model_options: List[MindFlowModel] = [
+            MindFlowModel.load(mindflow_model_id)
+            for mindflow_model_id in mindflow_model_ids
+        ]
+        mindflow_model_descriptions: List[str] = [
+            mindflow_model.name for mindflow_model in mindflow_model_options
+        ]
+
+        selected_mindflow_model: MindFlowModel = select_option(
+            "Select MindFlow model. Enter #",
+            mindflow_model_options,
+            mindflow_model_descriptions,
+        )
+        if selected_mindflow_model.id == MindFlowModelID.QUERY.value:
+            configure_query_model()
+        elif selected_mindflow_model.id == MindFlowModelID.INDEX.value:
+            configure_index_model()
+        elif selected_mindflow_model.id == MindFlowModelID.EMBEDDING.value:
+            configure_embedding_model()
+
+    def configure_query_model():
+        model_ids = [
+            ModelID.GPT_3_5_TURBO.value,
+            ModelID.GPT_4.value,
+        ]
+        model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
+        model_descriptions: List[str] = [
+            model.config_description for model in model_options
+        ]
+
+        selected_model: Model = select_option(
+            "Select chat model. Recommended GPT-4/Claude V1. Enter #",
+            model_options,
+            model_descriptions,
+        )
+        mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
+            f"{MindFlowModelID.QUERY.value}_config"
+        ) or MindFlowModelConfig(f"{MindFlowModelID.QUERY.value}_config")
+        mindflow_model_config.model = selected_model.id
+        mindflow_model_config.save()
+
+        print(f"Query Model: {selected_model.id} saved!")
+
+    def configure_index_model():
+        model_ids = [
+            ModelID.GPT_3_5_TURBO.value,
+            ModelID.GPT_4.value,
+        ]
+        model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
+        model_descriptions: List[str] = [
+            model.config_description for model in model_options
+        ]
+
+        selected_model: Model = select_option(
+            "Select chat model. Recommended GPT-3.5 Turbo/Claude Instant V1. Enter #",
+            model_options,
+            model_descriptions,
+        )
+        mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
+            f"{MindFlowModelID.INDEX.value}_config"
+        ) or MindFlowModelConfig(f"{MindFlowModelID.INDEX.value}_config")
+        mindflow_model_config.model = selected_model.id
+        mindflow_model_config.save()
+
+        print(f"Index Model: {selected_model.id} saved!")
+
+    def configure_embedding_model():
+        model_ids = [ModelID.TEXT_EMBEDDING_ADA_002.value]
+        model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
+        model_descriptions: List[str] = [
+            model.config_description for model in model_options
+        ]
+
+        selected_model: Model = select_option(
+            "Select chat model. Only one option... for now :) Enter #",
+            model_options,
+            model_descriptions,
+        )
+        mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
+            f"{MindFlowModelID.EMBEDDING.value}_config"
+        ) or MindFlowModelConfig(f"{MindFlowModelID.EMBEDDING.value}_config")
+        mindflow_model_config.model = selected_model.id
+        mindflow_model_config.save()
+
+        print(f"Embedding Model: {selected_model.id} saved!")
+
     config_options = ["model"]
     selected_config = select_option(
         "What do you want to configure? Enter #", config_options, config_options
@@ -29,112 +120,6 @@ def config():
     save_json_store()
 
 
-def configure_model():
-    mindflow_model_ids = [
-        MindFlowModelID.QUERY.value,
-        MindFlowModelID.INDEX.value,
-        MindFlowModelID.EMBEDDING.value,
-    ]
-    mindflow_model_options: List[MindFlowModel] = [
-        MindFlowModel.load(mindflow_model_id)
-        for mindflow_model_id in mindflow_model_ids
-    ]
-    mindflow_model_descriptions: List[str] = [
-        mindflow_model.name for mindflow_model in mindflow_model_options
-    ]
-
-    selected_mindflow_model: MindFlowModel = select_option(
-        "Select MindFlow model. Enter #",
-        mindflow_model_options,
-        mindflow_model_descriptions,
-    )
-    if selected_mindflow_model.id == MindFlowModelID.QUERY.value:
-        configure_query_model()
-    elif selected_mindflow_model.id == MindFlowModelID.INDEX.value:
-        configure_index_model()
-    elif selected_mindflow_model.id == MindFlowModelID.EMBEDDING.value:
-        configure_embedding_model()
-
-
-def configure_query_model():
-    model_ids = [
-        ModelID.GPT_3_5_TURBO.value,
-        ModelID.GPT_4.value,
-        ModelID.CLAUDE_INSTANT_V1.value,
-        ModelID.CLAUDE_V1.value,
-    ]
-    model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
-    model_descriptions: List[str] = [
-        model.config_description for model in model_options
-    ]
-
-    selected_model: Model = select_option(
-        "Select chat model. Recommended GPT-4/Claude V1. Enter #",
-        model_options,
-        model_descriptions,
-    )
-    mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
-        f"{MindFlowModelID.QUERY.value}_config"
-    ) or MindFlowModelConfig(f"{MindFlowModelID.QUERY.value}_config")
-    mindflow_model_config.model = selected_model.id
-    mindflow_model_config.save()
-
-    print(f"Query Model: {selected_model.id} saved!")
-
-
-def configure_index_model():
-    model_ids = [
-        ModelID.GPT_3_5_TURBO.value,
-        ModelID.GPT_4.value,
-        ModelID.CLAUDE_INSTANT_V1.value,
-        ModelID.CLAUDE_V1.value,
-    ]
-    model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
-    model_descriptions: List[str] = [
-        model.config_description for model in model_options
-    ]
-
-    selected_model: Model = select_option(
-        "Select chat model. Recommended GPT-3.5 Turbo/Claude Instant V1. Enter #",
-        model_options,
-        model_descriptions,
-    )
-    mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
-        f"{MindFlowModelID.INDEX.value}_config"
-    ) or MindFlowModelConfig(f"{MindFlowModelID.INDEX.value}_config")
-    mindflow_model_config.model = selected_model.id
-    mindflow_model_config.save()
-
-    print(f"Index Model: {selected_model.id} saved!")
-
-
-def configure_embedding_model():
-    model_ids = [ModelID.TEXT_EMBEDDING_ADA_002.value]
-    model_options: List[Model] = [Model.load(model_id) for model_id in model_ids]
-    model_descriptions: List[str] = [
-        model.config_description for model in model_options
-    ]
-
-    selected_model: Model = select_option(
-        "Select chat model. Only one option... for now :) Enter #",
-        model_options,
-        model_descriptions,
-    )
-    mindflow_model_config: MindFlowModelConfig = MindFlowModelConfig.load(
-        f"{MindFlowModelID.EMBEDDING.value}_config"
-    ) or MindFlowModelConfig(f"{MindFlowModelID.EMBEDDING.value}_config")
-    mindflow_model_config.model = selected_model.id
-    mindflow_model_config.save()
-
-    print(f"Embedding Model: {selected_model.id} saved!")
-
-
-def clear_console(lines: int):
-    for _ in range(lines):
-        sys.stdout.write("\033[F")  # Move cursor up one line
-        sys.stdout.write("\033[K")  # Clear the line
-
-
 def select_option(prompt: str, options: List, descriptions: List[str]) -> int:
     for i, description in enumerate(descriptions, 1):
         click.echo(f"{i}: {description}")
@@ -153,3 +138,11 @@ def select_option(prompt: str, options: List, descriptions: List[str]) -> int:
 
     clear_console(lines_to_clear)
     return options[selected_option_index - 1]
+
+
+def clear_console(lines: int):
+    import sys
+
+    for _ in range(lines):
+        sys.stdout.write("\033[F")  # Move cursor up one line
+        sys.stdout.write("\033[K")  # Clear the line