reverse-engineering-assistant 1.0.3__tar.gz → 2.9.2__tar.gz

{reverse-engineering-assistant-1.0.3 → reverse_engineering_assistant-2.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: reverse-engineering-assistant
-Version: 1.0.3
+Version: 2.9.2
 Summary: An AI assistant for reverse engineering tasks
 Author: サイバーカイダ (cyberkaida)
 Classifier: License :: OSI Approved :: Apache Software License
@@ -11,17 +11,19 @@ Description-Content-Type: text/markdown
 Requires-Dist: langchain
 Requires-Dist: langchain-core
 Requires-Dist: langchain-openai
-Requires-Dist: llama-cpp-python
+Requires-Dist: langchain-community
 Requires-Dist: prompt_toolkit
 Requires-Dist: sentence_transformers
 Requires-Dist: PyYAML
 Requires-Dist: pydantic
 Requires-Dist: rich
-Requires-Dist: Flask
+Requires-Dist: grpcio
+Requires-Dist: protobuf
+Requires-Dist: flask
 
 # ReVA - Reverse Engineering Assistant
 
-[✨ An (old) quick demo! ✨](https://asciinema.org/a/626197)
+> Updated demo coming soon!
 
 The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
 reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
@@ -46,13 +48,14 @@ information from the tools, but when there is no information it can still respon
 questions from its training.
 
 You can ask questions like:
+- What are the interesting strings in this program?
 - Does this program use encryption? Write a markdown report on the encryption and where it is used.
 - Draw a class diagram using plantuml syntax.
 - Start from main, examine the program in detail. Rename variables as you go and provide a summary of the program.
 - Explain the purpose of the `__mod_init` segment.
 - What does `mmap` return?
 - What does the function at address 0x80000 do?
-- This is a CTF problem. Start at main, examine the program in detail and write a pwntools script to get the flag.
+- This is a CTF problem. Write a pwntools script to get the flag.
 
 An important part of reverse engineering is the process. Many other tools simply ask a single question of the LLM,
 this means it is difficult to determine _why_ a thing happened. In ReVa we break all actions down into small parts
@@ -68,36 +71,40 @@ Built in support is provided for:
 - [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
 - [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
 
-Limited support is provided for:
-- [llama-cpp](https://llama-cpp-python.readthedocs.io/en/latest/) and any model it supports for local on-device inference
-- [text-generation-webui](https://github.com/oobabooga/text-generation-webui) and any model it supports for self-hosted remote inference
+See [Configuration](#configuration) for more information about settings for the providers.
 
 Adding additional inference servers is easy if it is supported by langchain.
 
-See the configuration section for more information about setting the model.
-
 ## Configuration
 
-Configuration for the reverse engineering assistant is stored at
-`~/.config/reverse-engineering-assistant/config.yaml`. If this
-is not present on first start, a default configuration using
-OpenAI for inference and the `OPENAI_API_TOKEN` environment
-variable will be used.
+Configuration for ReVa is in the CodeBrowser Tool options.
+Open a program and go to Edit -> Tool Options -> ReVa.
+
+There are options for:
+- Selecting a provider (OpenAI or Ollama, others coming soon!)
+- Enabling "Follow", this will move the Ghidra view to the location of
+things ReVa is examining or changing.
+- Enabling "Auto-allow", ReVa will log her actions for the user to accept
+in the "ReVa Actions Log" window.
+
+There are sections for the providers.
+
+### OpenAI
 
-The most important setting is the `type` top level setting.
-This controls what inference service you use. These are the
-same as the configuration keys, for example to use Ollama,
-set type to `ollama` and configure the settings in the `ollama:`
-section.
+By default, the OpenAI key is loaded from the environment variable `OPENAI_API_KEY`. You can also set your key inside Ghidra. Setting the key back to the `OPENAI_API_KEY` value will clear the key from the Ghidra configuration and load it from the environment.
 
-The configuration also contains the prompts used for the models.
-If you use Ollama or OpenAI these will be processed to fit the
-model specific prompt pattern (placing the system prompt in the
-correct tags, etc).
+You can also select the model. By default `gpt-4o` is selected. This model works best with the tools and the prompt provided by ReVa.
 
-For `llama-cpp` and `text-generation-webui` these may need to be
-configured for your specific model. For this reason Ollama is
-preferred for self hosting.
+`gpt-4` also works well, but is slow and needs more prompting by the user to explore a binary.
+
+### Ollama
+
+Ollama is a local inference server. The default server is set to localhost, with the default Ollama port. You can change this to a remote server if you want to perform inference on a remote machine. This is useful for organisations that self host.
+
+You can also select a model. The model must alread be loaded on the server. Good performance has been seen with:
+- `mixtral`
+- `llama3`
+- `phi`
 
 ## Workflow
 
@@ -106,54 +113,68 @@ RevA has a two step workflow.
 2. Open the chat session.
 
 ReVa uses an extension for your RE tool to perform analysis.
-See [Ghidra Support](#ghidra-support) and [BinaryNinja Support](#binary-ninja-support) below.
+See [Ghidra Support](#ghidra-support) below.
 
-Once open the RE tool will try to connect to ReVa's REST API on localhost.
+To ask questions and run the inference a command line tool is provided. Run `reva-chat` to begin the chat session. This command will find your open Ghidra
+and connect to it. To open a new chat, run the command again in another terminal.
 
-A project cache is created in `~/.cache/reverse-engineering-assistant/projects`. This contains your chat log and other
-cache data. This can be deleted at any time and ReVa will re-generate the data as needed.
+If you have more than one Ghidra open, you can select the right one with
+`reva-chat --project ${project-name}`, if it is not set, `reva-chat` will
+ask you which project you want to connect to.
 
-To ask questions and run the inference a command line tool is provided. Run `revassistant --project ${NAME_OF_YOUR_FILE}` to begin the chat session.
+## Protocol Build
 
-> Note: In the future `--project` will refer to a _project_ in Ghidra and allow inference across multiple files.
-> I am waiting for BinaryNinja's project feature to make this change, if this takes too long I will rework this argument.
+To communicate between `reva-server` and the extension, [gRPC](https://grpc.io) is used. You can read more about that (here)[./DEVELOPER.md]. Building the source files from those protocol definitions is driven from the [Makefile](/Makefile). To build the protocol source code files, run this command in the project's root:
 
-`revassistant` provides a chat window and runs the command API to talk with the RE tool.
+```sh
+make protocol
+```
 
-> Note: Right now only one `revassistant` can run at a time (as we start a server on a well known port)
-> In the future we will share the server between chat clients and RE tool connections.
+## Python Project (reva-server and reva-chat) Installation
 
-## Installation
+First install the python component, I like to use `pipx`. Install it with something like:
 
-To install the particular extension for your disassembler see:
-- [Ghidra Support](#ghidra-support)
-- [Binary Ninja Support](#binary-ninja-support)
+```sh
+pip install pipx
+```
 
-To install the chat component you can do the following:
+In the `reverse-engineering-assistant` folder, run:
 
 ```sh
-python3 -m pip install ./reverse-engineering-assistant
+pipx install .
 ```
 
-The chat can be started with:
+After installing the python project, pipx may warn you that you need to add a folder to your PATH environment variable. Make sure that the folder (now containing `reva-server` and `reva-chat`) are in your PATH variable. pipx can do it for you with this command:
 
 ```sh
-revassistant --project ${NAME_OF_YOUR_PROJECT}
+pipx ensurepath
+```
+
+The extension will need to start `reva-server`, and you will need to run `reva-chat`. In case you do not want to add them to your PATH, see the [Configuration](#configuration) section for how to set the path to the executables.
+
+Once the `reva-server` has been started by the extension the chat can be started with:
+
+```sh
+reva-chat
 ```
 
 # Ghidra Support
 
 ## Usage
 
-The [ghidra-assistant](ghidra-assistant/README.md) plugin must be installed first. 
+> The Python package must be installed for the Ghidra extension to work!
+
+Follow the instructions in the [ghidra-assistant](ghidra-assistant/README.md) plugin.
 
-After installation, enable the `ReVaPlugin` extension in the CodeBrowser tool (Open a file and click: File -> Configure -> Miscellaneous).
+After installation, enable the `ReVa Plugin` extension in the CodeBrowser tool (Open a file and click: File -> Configure -> Miscellaneous).
 
 If you want ReVa enabled by default, click File -> Save Tool to save the configuration.
 
-To start the inference side, open Help -> About ${program name}. In this popup you will see details about your open file.
-The `Program Name:` field is the name you need to pass to `revassistant --project` to start the inference server. In some
-cases this is different to the name in the project view.
+If everything is working correctly you will see a ReVa menu on your menu bar.
+
+## Configuration
+
+You can modify the plugin configuration in `Edit -> Tool Options -> ReVa`.
 
 ## Undo
 
@@ -162,28 +183,22 @@ one undo.
 
 ## Menus
 
-ReVa adds some elements to the Ghidra UI. You can either ask ReVa to do something in the chat window,
-"Examine the variable usage in `main` in detail, rename the variables with more descriptive names.",
-or use the menu system.
-
-For example you can right click a variable in the decompilation, select Reva -> Rename variable and ReVa
-will perform the action.
+ReVa adds an option to the CodeBrowser Tool's Window menu.
+Select Window -> ReVa Action Log to open the ReVa Action Log window.
 
-Note this uses the same system as chatting with ReVa, this means you can monitor ReVas thoughts in the chat
-window while the action is performed.
+This window shows actions ReVa has performed and would like to perform.
+You can accept or reject a change by double clicking the ✅ or ❌ icon. You can also go to the location the action will be performed by double clicking the address.
 
-# Binary Ninja Support
+If you reject an action, ReVa will be told and she will move on.
 
-> Note: Binary Ninja support is currently on hold while the basic functions are implemented in the Ghidra plugin.
-> This is because plugin development for Binary Ninja is easier as we have Python3. I will resume development soon!
+You can also enable "Auto-allow" in the ReVa options. This will automatically accept all actions ReVa wants to perform.
 
-Install the ReVA BinaryNinja plugin by opening your BinaryNinja plugin directory (Plugins -> Open Plugin Folder)
-and copying or symbolic linking the [binary-ninja-assistant](./binary-ninja-assistant) directory into the plugin
-directory.
+ReVa also adds some elements to the Ghidra UI. You can either ask ReVa to do something in the chat window,
+"Examine the variable usage in `main` in detail, rename the variables with more descriptive names.",
+or use the menu system.
 
-Restart Binary Ninja and "ReVA Push" will be available in the Plugin menu.
-Press this to push data from BinaryNinja to ReVA, then follow the instructions in the [Workflow section](#workflow).
-The project name will be the name of the current open file.
+For example you can right click a variable in the decompilation, select Reva -> Rename variable and ReVa
+will perform the action.
 
 # Support
 

reverse_engineering_assistant-2.9.2/README.md

@@ -0,0 +1,183 @@
+# ReVA - Reverse Engineering Assistant
+
+> Updated demo coming soon!
+
+The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
+reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
+
+ReVa is different from other efforts at building AI assistants for RE tasks because it uses a _tool driven approach_.
+ReVa aims to provide a variety of small tools to the LLM, just as your RE environment provides a set of small tools
+to you. ReVa combines this approach with chain-of-reasoning techniques to empower the LLM to complete complex tasks.
+
+Each of the tools given to the LLM are constructed to be easy for the LLM to use and to tolerate a variety of inputs
+and to reduce hallucination by the LLM. We do this by providing the LLM with a schema but tolerating other input,
+including descriptions that guide the LLM,and redirecting correctable mistakes back to the LLM, and including extra
+output to guide the next decision by the LLM.
+
+For example, when the LLM requests decompilation from your RE tool, we will accept a raw address in hex, a raw address
+in base 10, a symbol name with a namespace, or a symbol. If the LLM gives us bad input we report this to the LLM along
+with instructions to correct the input (maybe encouraging it to use the function list for example). To encourage exploration
+as a human would, we report additional context like the namespace and cross references along with the decompilation, this
+is a small nudge to make the LLM explore the binary in the same way a human would.
+
+Using this technique you can ask general questions and get relevant answers. The model prioritises
+information from the tools, but when there is no information it can still respond to generic
+questions from its training.
+
+You can ask questions like:
+- What are the interesting strings in this program?
+- Does this program use encryption? Write a markdown report on the encryption and where it is used.
+- Draw a class diagram using plantuml syntax.
+- Start from main, examine the program in detail. Rename variables as you go and provide a summary of the program.
+- Explain the purpose of the `__mod_init` segment.
+- What does `mmap` return?
+- What does the function at address 0x80000 do?
+- This is a CTF problem. Write a pwntools script to get the flag.
+
+An important part of reverse engineering is the process. Many other tools simply ask a single question of the LLM,
+this means it is difficult to determine _why_ a thing happened. In ReVa we break all actions down into small parts
+and include the LLMs thoughts in the output. This allows the analyst to monitor the LLMs actions and reasoning, aborting
+and changing the prompt if required.
+
+## Large Language Model Support
+
+RevA is based on [langchain](https://langchain.com),
+which supports a number of models.
+
+Built in support is provided for:
+- [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
+- [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
+
+See [Configuration](#configuration) for more information about settings for the providers.
+
+Adding additional inference servers is easy if it is supported by langchain.
+
+## Configuration
+
+Configuration for ReVa is in the CodeBrowser Tool options.
+Open a program and go to Edit -> Tool Options -> ReVa.
+
+There are options for:
+- Selecting a provider (OpenAI or Ollama, others coming soon!)
+- Enabling "Follow", this will move the Ghidra view to the location of
+things ReVa is examining or changing.
+- Enabling "Auto-allow", ReVa will log her actions for the user to accept
+in the "ReVa Actions Log" window.
+
+There are sections for the providers.
+
+### OpenAI
+
+By default, the OpenAI key is loaded from the environment variable `OPENAI_API_KEY`. You can also set your key inside Ghidra. Setting the key back to the `OPENAI_API_KEY` value will clear the key from the Ghidra configuration and load it from the environment.
+
+You can also select the model. By default `gpt-4o` is selected. This model works best with the tools and the prompt provided by ReVa.
+
+`gpt-4` also works well, but is slow and needs more prompting by the user to explore a binary.
+
+### Ollama
+
+Ollama is a local inference server. The default server is set to localhost, with the default Ollama port. You can change this to a remote server if you want to perform inference on a remote machine. This is useful for organisations that self host.
+
+You can also select a model. The model must alread be loaded on the server. Good performance has been seen with:
+- `mixtral`
+- `llama3`
+- `phi`
+
+## Workflow
+
+RevA has a two step workflow.
+1. Open your RE tool and the program you want to examine
+2. Open the chat session.
+
+ReVa uses an extension for your RE tool to perform analysis.
+See [Ghidra Support](#ghidra-support) below.
+
+To ask questions and run the inference a command line tool is provided. Run `reva-chat` to begin the chat session. This command will find your open Ghidra
+and connect to it. To open a new chat, run the command again in another terminal.
+
+If you have more than one Ghidra open, you can select the right one with
+`reva-chat --project ${project-name}`, if it is not set, `reva-chat` will
+ask you which project you want to connect to.
+
+## Protocol Build
+
+To communicate between `reva-server` and the extension, [gRPC](https://grpc.io) is used. You can read more about that (here)[./DEVELOPER.md]. Building the source files from those protocol definitions is driven from the [Makefile](/Makefile). To build the protocol source code files, run this command in the project's root:
+
+```sh
+make protocol
+```
+
+## Python Project (reva-server and reva-chat) Installation
+
+First install the python component, I like to use `pipx`. Install it with something like:
+
+```sh
+pip install pipx
+```
+
+In the `reverse-engineering-assistant` folder, run:
+
+```sh
+pipx install .
+```
+
+After installing the python project, pipx may warn you that you need to add a folder to your PATH environment variable. Make sure that the folder (now containing `reva-server` and `reva-chat`) are in your PATH variable. pipx can do it for you with this command:
+
+```sh
+pipx ensurepath
+```
+
+The extension will need to start `reva-server`, and you will need to run `reva-chat`. In case you do not want to add them to your PATH, see the [Configuration](#configuration) section for how to set the path to the executables.
+
+Once the `reva-server` has been started by the extension the chat can be started with:
+
+```sh
+reva-chat
+```
+
+# Ghidra Support
+
+## Usage
+
+> The Python package must be installed for the Ghidra extension to work!
+
+Follow the instructions in the [ghidra-assistant](ghidra-assistant/README.md) plugin.
+
+After installation, enable the `ReVa Plugin` extension in the CodeBrowser tool (Open a file and click: File -> Configure -> Miscellaneous).
+
+If you want ReVa enabled by default, click File -> Save Tool to save the configuration.
+
+If everything is working correctly you will see a ReVa menu on your menu bar.
+
+## Configuration
+
+You can modify the plugin configuration in `Edit -> Tool Options -> ReVa`.
+
+## Undo
+
+Whenever ReVa performs an action it will create an undo point for each action. If ReVa renames 5 variables, this will be
+one undo.
+
+## Menus
+
+ReVa adds an option to the CodeBrowser Tool's Window menu.
+Select Window -> ReVa Action Log to open the ReVa Action Log window.
+
+This window shows actions ReVa has performed and would like to perform.
+You can accept or reject a change by double clicking the ✅ or ❌ icon. You can also go to the location the action will be performed by double clicking the address.
+
+If you reject an action, ReVa will be told and she will move on.
+
+You can also enable "Auto-allow" in the ReVa options. This will automatically accept all actions ReVa wants to perform.
+
+ReVa also adds some elements to the Ghidra UI. You can either ask ReVa to do something in the chat window,
+"Examine the variable usage in `main` in detail, rename the variables with more descriptive names.",
+or use the menu system.
+
+For example you can right click a variable in the decompilation, select Reva -> Rename variable and ReVa
+will perform the action.
+
+# Support
+
+Do you like my work? Want to support this project and others? Interested in how this project was designed and built?
+This project and many others are built live on my stream at https://twitch.tv/cyberkaida !

{reverse-engineering-assistant-1.0.3 → reverse_engineering_assistant-2.9.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "reverse-engineering-assistant"
 readme = "README.md"
-version = "1.0.3"
+version = "2.9.2"
 authors = [
     {name="サイバーカイダ (cyberkaida)"},
 ]
@@ -20,17 +20,17 @@ dependencies = [
     "langchain",
     "langchain-core",
     "langchain-openai",
-    "llama-cpp-python",
+    "langchain-community",
     "prompt_toolkit",
     "sentence_transformers",
     "PyYAML",
     "pydantic",
     "rich",
-    "Flask",
+    "grpcio",
+    "protobuf",
+    "flask",
 ]
 
 [project.scripts]
-revassistant = "reverse_engineering_assistant.assistant:main"
-reva-serve = "reverse_engineering_assistant.assistant_api_server:main"
 reva-server = "reverse_engineering_assistant.assistant_api_server:main"
-
+reva-chat = "reverse_engineering_assistant.chat_client:main"

reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/.gitignore

@@ -0,0 +1,2 @@
+protocol/
+*_pb2.*

reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/api_server_tools/__init__.py

@@ -0,0 +1,16 @@
+
+
+from typing import List
+
+
+class RevaMessageHandler(object):
+    handles_type = None
+
+
+_global_message_handlers: List[RevaMessageHandler] = []
+
+__all__ = ['register_message_handler']
+
+def register_message_handler(cls: RevaMessageHandler):
+    _global_message_handlers.append(cls)
+    return cls

reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/api_server_tools/connection.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+import sys
+sys.path.append(str(Path(__file__).parent.joinpath('protocol')))
+import grpc
+from grpc import Channel, Server
+
+from typing import Optional
+
+from functools import cache
+
+_channel: Optional[Channel] = None
+@cache
+def get_channel() -> Channel:
+    global _channel
+    if not _channel:
+        raise ValueError("Channel not set")
+    return _channel
+
+@cache
+def connect_to_extension(host: str, port: int) -> Channel:
+    channel: Channel = grpc.insecure_channel(f"{host}:{port}")
+    global _channel
+    _channel = channel
+    return get_channel()

reverse_engineering_assistant-2.9.2/reverse_engineering_assistant/api_server_tools/llm_tools.py

@@ -0,0 +1,159 @@
+
+
+from ast import Call
+import queue
+import threading
+from typing import Callable, Dict, Optional
+from venv import logger
+
+from ..protocol.RevaChat_pb2_grpc import RevaChatServiceServicer
+from ..protocol.RevaChat_pb2 import RevaChatMessageResponse
+
+from functools import cache
+from ..assistant import ReverseEngineeringAssistant
+
+import logging
+module_logger = logging.getLogger("reva-server")
+
+from langchain_core.callbacks.base import BaseCallbackHandler
+from langchain_core.agents import AgentAction, AgentFinish
+
+from langchain_core.language_models.base import BaseLanguageModel
+from langchain_core.language_models.chat_models import BaseChatModel
+
+class RevaActionCollector(BaseCallbackHandler):
+    """
+    A callback handler for logging agent actions in the reverse engineering assistant.
+
+    This class logs agent actions and calls a callback. This is what prints the green
+    thoughts from the model to the console. This is very useful for the analyst to understand
+    what the model is doing (and is arguably the most important part of the assistant output!)
+
+    Attributes:
+        callback (Callable[[str], None]): The callback function to call when an agent action is performed.
+        logger (logging.Logger): The logger instance for the reverse_engineering_assistant.RevaActionLogger class.
+    """
+
+    callback: Callable[[str], None]
+    def __init__(self, callback: Callable[[str], None]) -> None:
+        super().__init__()
+        self.callback = callback
+
+    logger = logging.getLogger("reverse_engineering_assistant.RevaActionLogger")
+
+    def on_agent_action(self, action: AgentAction, **kwargs) -> None:
+        """
+        Callback method called when an agent action is performed.
+
+        Args:
+            action (AgentAction): The agent action that was performed.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            None
+        """
+        logger.debug(f"Agent action: {action} {kwargs}")
+        # TODO: Should this be AgentAction?
+        # TODO: Is `.action` still a thing?
+        self.callback(str(action.log))
+
+
+class RevaChat(RevaChatServiceServicer):
+    logger = logging.getLogger("reva-server.RevaChat")
+    llm: BaseLanguageModel | BaseChatModel
+
+    def __init__(self, llm: BaseLanguageModel | BaseChatModel) -> None:
+        self.llm = llm
+
+    def chat(self, request, context):
+        self.logger.info(f"Received request: {request}")
+        assistant = ReverseEngineeringAssistant(
+            request.project,
+            model=self.llm
+        )
+        self.logger.info(f"Assistant: {assistant}")
+        llm_response = assistant.query(request.message)
+        self.logger.info(f"LLM Response: {llm_response}")
+        response = RevaChatMessageResponse()
+        response.message = llm_response
+        return response
+
+    def chatResponseStream(self, request, context):
+        """
+        Given a request, return a stream of responses including
+        thoughts and a final message from the LLM.
+        """
+        self.logger.info(f"Received request: {request}")
+
+        response_queue: queue.Queue = queue.Queue()
+
+        def callback(message: str):
+            # Called for intermediate thoughts
+            response = RevaChatMessageResponse()
+            response.thought = message
+            response_queue.put(response)
+
+        assistant = ReverseEngineeringAssistant(
+            request.project,
+            model=self.llm,
+            langchain_callbacks=[RevaActionCollector(callback)]
+        )
+        self.logger.info(f"Assistant: {assistant}")
+
+        def run_query(query: str):
+            llm_response = assistant.query(query)
+            self.logger.info(f"LLM Response: {llm_response}")
+            response = RevaChatMessageResponse()
+            response.message = llm_response
+            response_queue.put(response)
+
+        t = threading.Thread(target=run_query, args=[request.message])
+        t.start()
+
+        done = False
+        while not done:
+            response = response_queue.get()
+            # We stop when we get a message and not thoughts
+            if response.message and not response.thought:
+                done = True
+            yield response
+        t.join()
+
+    def chatStream(self, request_iterator, context):
+        assistant: Optional[ReverseEngineeringAssistant] = None
+        response_queue: queue.Queue = queue.Queue()
+
+        def callback(message: str):
+            # Called for intermediate thoughts
+            response = RevaChatMessageResponse()
+            response.thought = message
+            response_queue.put(response)
+
+        for request in request_iterator:
+            if not assistant:
+                assistant = ReverseEngineeringAssistant(
+                    request.project,
+                    model=self.llm,
+                    langchain_callbacks=[RevaActionCollector(callback)]
+                )
+            self.logger.info(f"Received request: {request}")
+            def run_query(query: str):
+                assert assistant is not None
+                self.logger.info(f"Asking assistant: {query}")
+                llm_response = assistant.query(query)
+                self.logger.info(f"LLM Response: {llm_response}")
+                response = RevaChatMessageResponse()
+                response.message = llm_response
+                response_queue.put(response)
+
+            t = threading.Thread(target=run_query, args=[request.message])
+            t.start()
+
+            done = False
+            while not done:
+                response = response_queue.get()
+                # We stop when we get a message and not thoughts
+                if response.message and not response.thought:
+                    done = True
+                yield response
+            t.join()