reverse-engineering-assistant 1.0.1__tar.gz → 2.9.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/PKG-INFO +29 -51
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/README.md +26 -48
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/pyproject.toml +4 -6
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/.gitignore +2 -0
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/__init__.py +16 -0
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/connection.py +24 -0
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/llm_tools.py +111 -0
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/re_tools.py +354 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/assistant.py +6 -1
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/assistant_api_server.py +111 -0
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/chat_client.py +179 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant.egg-info/PKG-INFO +29 -51
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant.egg-info/SOURCES.txt +5 -3
- reverse_engineering_assistant-2.9.1/reverse_engineering_assistant.egg-info/entry_points.txt +3 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant.egg-info/requires.txt +2 -2
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant/api_server_tools/__init__.py +0 -0
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant/api_server_tools/function_tools.py +0 -277
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant/api_server_tools/llm_tools.py +0 -39
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant/assistant_api_server.py +0 -355
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant/tool_protocol.py +0 -460
- reverse-engineering-assistant-1.0.1/reverse_engineering_assistant.egg-info/entry_points.txt +0 -4
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/__init__.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/configuration.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/documents.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/model.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/reva_exceptions.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant/tool.py +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant.egg-info/dependency_links.txt +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/reverse_engineering_assistant.egg-info/top_level.txt +0 -0
- {reverse-engineering-assistant-1.0.1 → reverse_engineering_assistant-2.9.1}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: reverse-engineering-assistant
|
|
3
|
-
Version:
|
|
3
|
+
Version: 2.9.1
|
|
4
4
|
Summary: An AI assistant for reverse engineering tasks
|
|
5
5
|
Author: サイバーカイダ (cyberkaida)
|
|
6
6
|
Classifier: License :: OSI Approved :: Apache Software License
|
|
@@ -11,17 +11,17 @@ Description-Content-Type: text/markdown
|
|
|
11
11
|
Requires-Dist: langchain
|
|
12
12
|
Requires-Dist: langchain-core
|
|
13
13
|
Requires-Dist: langchain-openai
|
|
14
|
-
Requires-Dist: llama-cpp-python
|
|
15
14
|
Requires-Dist: prompt_toolkit
|
|
16
15
|
Requires-Dist: sentence_transformers
|
|
17
16
|
Requires-Dist: PyYAML
|
|
18
17
|
Requires-Dist: pydantic
|
|
19
18
|
Requires-Dist: rich
|
|
20
|
-
Requires-Dist:
|
|
19
|
+
Requires-Dist: grpcio
|
|
20
|
+
Requires-Dist: protobuf
|
|
21
21
|
|
|
22
22
|
# ReVA - Reverse Engineering Assistant
|
|
23
23
|
|
|
24
|
-
|
|
24
|
+
> Updated demo coming soon!
|
|
25
25
|
|
|
26
26
|
The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
|
|
27
27
|
reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
|
|
@@ -68,16 +68,13 @@ Built in support is provided for:
|
|
|
68
68
|
- [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
|
|
69
69
|
- [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
|
|
70
70
|
|
|
71
|
-
Limited support is provided for:
|
|
72
|
-
- [llama-cpp](https://llama-cpp-python.readthedocs.io/en/latest/) and any model it supports for local on-device inference
|
|
73
|
-
- [text-generation-webui](https://github.com/oobabooga/text-generation-webui) and any model it supports for self-hosted remote inference
|
|
74
|
-
|
|
75
71
|
Adding additional inference servers is easy if it is supported by langchain.
|
|
76
72
|
|
|
77
|
-
See the configuration section for more information about setting the model.
|
|
78
|
-
|
|
79
73
|
## Configuration
|
|
80
74
|
|
|
75
|
+
> This is currently being moved to the Ghidra GUI
|
|
76
|
+
> See Edit -> Tool Options -> ReVa in the Codebrowser Tool
|
|
77
|
+
|
|
81
78
|
Configuration for the reverse engineering assistant is stored at
|
|
82
79
|
`~/.config/reverse-engineering-assistant/config.yaml`. If this
|
|
83
80
|
is not present on first start, a default configuration using
|
|
@@ -106,52 +103,47 @@ RevA has a two step workflow.
|
|
|
106
103
|
2. Open the chat session.
|
|
107
104
|
|
|
108
105
|
ReVa uses an extension for your RE tool to perform analysis.
|
|
109
|
-
See [Ghidra Support](#ghidra-support)
|
|
110
|
-
|
|
111
|
-
Once open the RE tool will try to connect to ReVa's REST API on localhost.
|
|
112
|
-
|
|
113
|
-
A project cache is created in `~/.cache/reverse-engineering-assistant/projects`. This contains your chat log and other
|
|
114
|
-
cache data. This can be deleted at any time and ReVa will re-generate the data as needed.
|
|
115
|
-
|
|
116
|
-
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.
|
|
106
|
+
See [Ghidra Support](#ghidra-support) below.
|
|
117
107
|
|
|
118
|
-
|
|
119
|
-
|
|
108
|
+
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
|
|
109
|
+
and connect to it. To open a new chat, run the command again in another terminal.
|
|
120
110
|
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
> In the future we will share the server between chat clients and RE tool connections.
|
|
111
|
+
If you have more than one Ghidra open, you can select the right one with
|
|
112
|
+
`reva-chat --project ${project-name}`, if it is not set, `reva-chat` will
|
|
113
|
+
ask you which project you want to connect to.
|
|
125
114
|
|
|
126
115
|
## Installation
|
|
127
116
|
|
|
117
|
+
First install the python component, I like to use `pipx`. It is best to make
|
|
118
|
+
sure that `reva-server` and `reva-chat` are on your path.
|
|
119
|
+
The Ghidra extension will need to start `reva-server`, and you will need to
|
|
120
|
+
run `reva-chat`.
|
|
121
|
+
|
|
128
122
|
To install the particular extension for your disassembler see:
|
|
129
123
|
- [Ghidra Support](#ghidra-support)
|
|
130
|
-
- [Binary Ninja Support](#binary-ninja-support)
|
|
131
|
-
|
|
132
|
-
To install the chat component you can do the following:
|
|
133
|
-
|
|
134
|
-
```sh
|
|
135
|
-
python3 -m pip install ./reverse-engineering-assistant
|
|
136
|
-
```
|
|
137
124
|
|
|
138
125
|
The chat can be started with:
|
|
139
126
|
|
|
140
127
|
```sh
|
|
141
|
-
|
|
128
|
+
reva-chat
|
|
142
129
|
```
|
|
143
130
|
|
|
131
|
+
> You can also configure the path to `reva-server` in `Edit -> Tool Options -> ReVa`
|
|
132
|
+
> if it is not on your path. But you really should put it on your path!
|
|
133
|
+
|
|
144
134
|
# Ghidra Support
|
|
145
135
|
|
|
146
136
|
## Usage
|
|
147
137
|
|
|
138
|
+
> The Python package must be installed for the Ghidra extension to work!
|
|
139
|
+
|
|
140
|
+
Follow the instructions in the [ghidra-assistant](ghidra-assistant/README.md) plugin.
|
|
141
|
+
|
|
148
142
|
After installation, enable the `ReVaPlugin` extension in the CodeBrowser tool (Open a file and click: File -> Configure -> Miscellaneous).
|
|
149
143
|
|
|
150
144
|
If you want ReVa enabled by default, click File -> Save Tool to save the configuration.
|
|
151
145
|
|
|
152
|
-
|
|
153
|
-
The `Program Name:` field is the name you need to pass to `revassistant --project` to start the inference server. In some
|
|
154
|
-
cases this is different to the name in the project view.
|
|
146
|
+
If everything is working correctly you will see a ReVa menu on your menu bar.
|
|
155
147
|
|
|
156
148
|
## Undo
|
|
157
149
|
|
|
@@ -160,6 +152,8 @@ one undo.
|
|
|
160
152
|
|
|
161
153
|
## Menus
|
|
162
154
|
|
|
155
|
+
> These are being added in the next release
|
|
156
|
+
|
|
163
157
|
ReVa adds some elements to the Ghidra UI. You can either ask ReVa to do something in the chat window,
|
|
164
158
|
"Examine the variable usage in `main` in detail, rename the variables with more descriptive names.",
|
|
165
159
|
or use the menu system.
|
|
@@ -167,22 +161,6 @@ or use the menu system.
|
|
|
167
161
|
For example you can right click a variable in the decompilation, select Reva -> Rename variable and ReVa
|
|
168
162
|
will perform the action.
|
|
169
163
|
|
|
170
|
-
Note this uses the same system as chatting with ReVa, this means you can monitor ReVas thoughts in the chat
|
|
171
|
-
window while the action is performed.
|
|
172
|
-
|
|
173
|
-
# Binary Ninja Support
|
|
174
|
-
|
|
175
|
-
> Note: Binary Ninja support is currently on hold while the basic functions are implemented in the Ghidra plugin.
|
|
176
|
-
> This is because plugin development for Binary Ninja is easier as we have Python3. I will resume development soon!
|
|
177
|
-
|
|
178
|
-
Install the ReVA BinaryNinja plugin by opening your BinaryNinja plugin directory (Plugins -> Open Plugin Folder)
|
|
179
|
-
and copying or symbolic linking the [binary-ninja-assistant](./binary-ninja-assistant) directory into the plugin
|
|
180
|
-
directory.
|
|
181
|
-
|
|
182
|
-
Restart Binary Ninja and "ReVA Push" will be available in the Plugin menu.
|
|
183
|
-
Press this to push data from BinaryNinja to ReVA, then follow the instructions in the [Workflow section](#workflow).
|
|
184
|
-
The project name will be the name of the current open file.
|
|
185
|
-
|
|
186
164
|
# Support
|
|
187
165
|
|
|
188
166
|
Do you like my work? Want to support this project and others? Interested in how this project was designed and built?
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
# ReVA - Reverse Engineering Assistant
|
|
2
2
|
|
|
3
|
-
|
|
3
|
+
> Updated demo coming soon!
|
|
4
4
|
|
|
5
5
|
The reverse engineering assistant (ReVA) is a project to build a disassembler agnostic AI assistant for
|
|
6
6
|
reverse engineering tasks. This includes both _offline_ and online inference and a simple architecture.
|
|
@@ -47,16 +47,13 @@ Built in support is provided for:
|
|
|
47
47
|
- [OpenAI](https://platform.openai.com/overview) for online inference and easy setup (Needs an OpenAI API key)
|
|
48
48
|
- [Ollama](https://ollama.ai) and any model it supports for local on-device inference or connecting to a self hosted remote inference server.
|
|
49
49
|
|
|
50
|
-
Limited support is provided for:
|
|
51
|
-
- [llama-cpp](https://llama-cpp-python.readthedocs.io/en/latest/) and any model it supports for local on-device inference
|
|
52
|
-
- [text-generation-webui](https://github.com/oobabooga/text-generation-webui) and any model it supports for self-hosted remote inference
|
|
53
|
-
|
|
54
50
|
Adding additional inference servers is easy if it is supported by langchain.
|
|
55
51
|
|
|
56
|
-
See the configuration section for more information about setting the model.
|
|
57
|
-
|
|
58
52
|
## Configuration
|
|
59
53
|
|
|
54
|
+
> This is currently being moved to the Ghidra GUI
|
|
55
|
+
> See Edit -> Tool Options -> ReVa in the Codebrowser Tool
|
|
56
|
+
|
|
60
57
|
Configuration for the reverse engineering assistant is stored at
|
|
61
58
|
`~/.config/reverse-engineering-assistant/config.yaml`. If this
|
|
62
59
|
is not present on first start, a default configuration using
|
|
@@ -85,52 +82,47 @@ RevA has a two step workflow.
|
|
|
85
82
|
2. Open the chat session.
|
|
86
83
|
|
|
87
84
|
ReVa uses an extension for your RE tool to perform analysis.
|
|
88
|
-
See [Ghidra Support](#ghidra-support)
|
|
89
|
-
|
|
90
|
-
Once open the RE tool will try to connect to ReVa's REST API on localhost.
|
|
91
|
-
|
|
92
|
-
A project cache is created in `~/.cache/reverse-engineering-assistant/projects`. This contains your chat log and other
|
|
93
|
-
cache data. This can be deleted at any time and ReVa will re-generate the data as needed.
|
|
94
|
-
|
|
95
|
-
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.
|
|
85
|
+
See [Ghidra Support](#ghidra-support) below.
|
|
96
86
|
|
|
97
|
-
|
|
98
|
-
|
|
87
|
+
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
|
|
88
|
+
and connect to it. To open a new chat, run the command again in another terminal.
|
|
99
89
|
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
> In the future we will share the server between chat clients and RE tool connections.
|
|
90
|
+
If you have more than one Ghidra open, you can select the right one with
|
|
91
|
+
`reva-chat --project ${project-name}`, if it is not set, `reva-chat` will
|
|
92
|
+
ask you which project you want to connect to.
|
|
104
93
|
|
|
105
94
|
## Installation
|
|
106
95
|
|
|
96
|
+
First install the python component, I like to use `pipx`. It is best to make
|
|
97
|
+
sure that `reva-server` and `reva-chat` are on your path.
|
|
98
|
+
The Ghidra extension will need to start `reva-server`, and you will need to
|
|
99
|
+
run `reva-chat`.
|
|
100
|
+
|
|
107
101
|
To install the particular extension for your disassembler see:
|
|
108
102
|
- [Ghidra Support](#ghidra-support)
|
|
109
|
-
- [Binary Ninja Support](#binary-ninja-support)
|
|
110
|
-
|
|
111
|
-
To install the chat component you can do the following:
|
|
112
|
-
|
|
113
|
-
```sh
|
|
114
|
-
python3 -m pip install ./reverse-engineering-assistant
|
|
115
|
-
```
|
|
116
103
|
|
|
117
104
|
The chat can be started with:
|
|
118
105
|
|
|
119
106
|
```sh
|
|
120
|
-
|
|
107
|
+
reva-chat
|
|
121
108
|
```
|
|
122
109
|
|
|
110
|
+
> You can also configure the path to `reva-server` in `Edit -> Tool Options -> ReVa`
|
|
111
|
+
> if it is not on your path. But you really should put it on your path!
|
|
112
|
+
|
|
123
113
|
# Ghidra Support
|
|
124
114
|
|
|
125
115
|
## Usage
|
|
126
116
|
|
|
117
|
+
> The Python package must be installed for the Ghidra extension to work!
|
|
118
|
+
|
|
119
|
+
Follow the instructions in the [ghidra-assistant](ghidra-assistant/README.md) plugin.
|
|
120
|
+
|
|
127
121
|
After installation, enable the `ReVaPlugin` extension in the CodeBrowser tool (Open a file and click: File -> Configure -> Miscellaneous).
|
|
128
122
|
|
|
129
123
|
If you want ReVa enabled by default, click File -> Save Tool to save the configuration.
|
|
130
124
|
|
|
131
|
-
|
|
132
|
-
The `Program Name:` field is the name you need to pass to `revassistant --project` to start the inference server. In some
|
|
133
|
-
cases this is different to the name in the project view.
|
|
125
|
+
If everything is working correctly you will see a ReVa menu on your menu bar.
|
|
134
126
|
|
|
135
127
|
## Undo
|
|
136
128
|
|
|
@@ -139,6 +131,8 @@ one undo.
|
|
|
139
131
|
|
|
140
132
|
## Menus
|
|
141
133
|
|
|
134
|
+
> These are being added in the next release
|
|
135
|
+
|
|
142
136
|
ReVa adds some elements to the Ghidra UI. You can either ask ReVa to do something in the chat window,
|
|
143
137
|
"Examine the variable usage in `main` in detail, rename the variables with more descriptive names.",
|
|
144
138
|
or use the menu system.
|
|
@@ -146,22 +140,6 @@ or use the menu system.
|
|
|
146
140
|
For example you can right click a variable in the decompilation, select Reva -> Rename variable and ReVa
|
|
147
141
|
will perform the action.
|
|
148
142
|
|
|
149
|
-
Note this uses the same system as chatting with ReVa, this means you can monitor ReVas thoughts in the chat
|
|
150
|
-
window while the action is performed.
|
|
151
|
-
|
|
152
|
-
# Binary Ninja Support
|
|
153
|
-
|
|
154
|
-
> Note: Binary Ninja support is currently on hold while the basic functions are implemented in the Ghidra plugin.
|
|
155
|
-
> This is because plugin development for Binary Ninja is easier as we have Python3. I will resume development soon!
|
|
156
|
-
|
|
157
|
-
Install the ReVA BinaryNinja plugin by opening your BinaryNinja plugin directory (Plugins -> Open Plugin Folder)
|
|
158
|
-
and copying or symbolic linking the [binary-ninja-assistant](./binary-ninja-assistant) directory into the plugin
|
|
159
|
-
directory.
|
|
160
|
-
|
|
161
|
-
Restart Binary Ninja and "ReVA Push" will be available in the Plugin menu.
|
|
162
|
-
Press this to push data from BinaryNinja to ReVA, then follow the instructions in the [Workflow section](#workflow).
|
|
163
|
-
The project name will be the name of the current open file.
|
|
164
|
-
|
|
165
143
|
# Support
|
|
166
144
|
|
|
167
145
|
Do you like my work? Want to support this project and others? Interested in how this project was designed and built?
|
|
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "reverse-engineering-assistant"
|
|
7
7
|
readme = "README.md"
|
|
8
|
-
version = "
|
|
8
|
+
version = "2.9.1"
|
|
9
9
|
authors = [
|
|
10
10
|
{name="サイバーカイダ (cyberkaida)"},
|
|
11
11
|
]
|
|
@@ -20,17 +20,15 @@ dependencies = [
|
|
|
20
20
|
"langchain",
|
|
21
21
|
"langchain-core",
|
|
22
22
|
"langchain-openai",
|
|
23
|
-
"llama-cpp-python",
|
|
24
23
|
"prompt_toolkit",
|
|
25
24
|
"sentence_transformers",
|
|
26
25
|
"PyYAML",
|
|
27
26
|
"pydantic",
|
|
28
27
|
"rich",
|
|
29
|
-
"
|
|
28
|
+
"grpcio",
|
|
29
|
+
"protobuf",
|
|
30
30
|
]
|
|
31
31
|
|
|
32
32
|
[project.scripts]
|
|
33
|
-
revassistant = "reverse_engineering_assistant.assistant:main"
|
|
34
|
-
reva-serve = "reverse_engineering_assistant.assistant_api_server:main"
|
|
35
33
|
reva-server = "reverse_engineering_assistant.assistant_api_server:main"
|
|
36
|
-
|
|
34
|
+
reva-chat = "reverse_engineering_assistant.chat_client:main"
|
reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/__init__.py
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
|
|
2
|
+
|
|
3
|
+
from typing import List
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class RevaMessageHandler(object):
|
|
7
|
+
handles_type = None
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
_global_message_handlers: List[RevaMessageHandler] = []
|
|
11
|
+
|
|
12
|
+
__all__ = ['register_message_handler']
|
|
13
|
+
|
|
14
|
+
def register_message_handler(cls: RevaMessageHandler):
|
|
15
|
+
_global_message_handlers.append(cls)
|
|
16
|
+
return cls
|
reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/connection.py
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
import sys
|
|
3
|
+
sys.path.append(str(Path(__file__).parent.joinpath('protocol')))
|
|
4
|
+
import grpc
|
|
5
|
+
from grpc import Channel, Server
|
|
6
|
+
|
|
7
|
+
from typing import Optional
|
|
8
|
+
|
|
9
|
+
from functools import cache
|
|
10
|
+
|
|
11
|
+
_channel: Optional[Channel] = None
|
|
12
|
+
@cache
|
|
13
|
+
def get_channel() -> Channel:
|
|
14
|
+
global _channel
|
|
15
|
+
if not _channel:
|
|
16
|
+
raise ValueError("Channel not set")
|
|
17
|
+
return _channel
|
|
18
|
+
|
|
19
|
+
@cache
|
|
20
|
+
def connect_to_extension(host: str, port: int) -> Channel:
|
|
21
|
+
channel: Channel = grpc.insecure_channel(f"{host}:{port}")
|
|
22
|
+
global _channel
|
|
23
|
+
_channel = channel
|
|
24
|
+
return get_channel()
|
reverse_engineering_assistant-2.9.1/reverse_engineering_assistant/api_server_tools/llm_tools.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
|
|
2
|
+
|
|
3
|
+
from ast import Call
|
|
4
|
+
import queue
|
|
5
|
+
import threading
|
|
6
|
+
from typing import Callable
|
|
7
|
+
from venv import logger
|
|
8
|
+
|
|
9
|
+
from ..protocol.RevaChat_pb2_grpc import RevaChatServiceServicer
|
|
10
|
+
from ..protocol.RevaChat_pb2 import RevaChatMessageResponse
|
|
11
|
+
|
|
12
|
+
from functools import cache
|
|
13
|
+
from ..assistant import ReverseEngineeringAssistant
|
|
14
|
+
|
|
15
|
+
import logging
|
|
16
|
+
module_logger = logging.getLogger("reva-server")
|
|
17
|
+
|
|
18
|
+
from langchain_core.callbacks.base import BaseCallbackHandler
|
|
19
|
+
from langchain_core.agents import AgentAction, AgentFinish
|
|
20
|
+
|
|
21
|
+
class RevaActionCollector(BaseCallbackHandler):
|
|
22
|
+
"""
|
|
23
|
+
A callback handler for logging agent actions in the reverse engineering assistant.
|
|
24
|
+
|
|
25
|
+
This class logs agent actions and calls a callback. This is what prints the green
|
|
26
|
+
thoughts from the model to the console. This is very useful for the analyst to understand
|
|
27
|
+
what the model is doing (and is arguably the most important part of the assistant output!)
|
|
28
|
+
|
|
29
|
+
Attributes:
|
|
30
|
+
callback (Callable[[str], None]): The callback function to call when an agent action is performed.
|
|
31
|
+
logger (logging.Logger): The logger instance for the reverse_engineering_assistant.RevaActionLogger class.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
callback: Callable[[str], None]
|
|
35
|
+
def __init__(self, callback: Callable[[str], None]) -> None:
|
|
36
|
+
super().__init__()
|
|
37
|
+
self.callback = callback
|
|
38
|
+
|
|
39
|
+
logger = logging.getLogger("reverse_engineering_assistant.RevaActionLogger")
|
|
40
|
+
|
|
41
|
+
def on_agent_action(self, action: AgentAction, **kwargs) -> None:
|
|
42
|
+
"""
|
|
43
|
+
Callback method called when an agent action is performed.
|
|
44
|
+
|
|
45
|
+
Args:
|
|
46
|
+
action (AgentAction): The agent action that was performed.
|
|
47
|
+
**kwargs: Additional keyword arguments.
|
|
48
|
+
|
|
49
|
+
Returns:
|
|
50
|
+
None
|
|
51
|
+
"""
|
|
52
|
+
logger.debug(f"Agent action: {action} {kwargs}")
|
|
53
|
+
# TODO: Should this be AgentAction?
|
|
54
|
+
# TODO: Is `.action` still a thing?
|
|
55
|
+
self.callback(str(action.log))
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class RevaChat(RevaChatServiceServicer):
|
|
59
|
+
logger = logging.getLogger("reva-server.RevaChat")
|
|
60
|
+
|
|
61
|
+
def chat(self, request, context):
|
|
62
|
+
self.logger.info(f"Received request: {request}")
|
|
63
|
+
assistant = ReverseEngineeringAssistant(request.project, langchain_callbacks=[RevaActionCollector(callback)])
|
|
64
|
+
self.logger.info(f"Assistant: {assistant}")
|
|
65
|
+
llm_response = assistant.query(request.message)
|
|
66
|
+
self.logger.info(f"LLM Response: {llm_response}")
|
|
67
|
+
response = RevaChatMessageResponse()
|
|
68
|
+
response.message = llm_response
|
|
69
|
+
return response
|
|
70
|
+
|
|
71
|
+
def chatResponseStream(self, request, context):
|
|
72
|
+
"""
|
|
73
|
+
Given a request, return a stream of responses including
|
|
74
|
+
thoughts and a final message from the LLM.
|
|
75
|
+
"""
|
|
76
|
+
self.logger.info(f"Received request: {request}")
|
|
77
|
+
|
|
78
|
+
response_queue: queue.Queue = queue.Queue()
|
|
79
|
+
|
|
80
|
+
def callback(message: str):
|
|
81
|
+
# Called for intermediate thoughts
|
|
82
|
+
response = RevaChatMessageResponse()
|
|
83
|
+
response.thought = message
|
|
84
|
+
response_queue.put(response)
|
|
85
|
+
|
|
86
|
+
assistant = ReverseEngineeringAssistant(request.project, langchain_callbacks=[RevaActionCollector(callback)])
|
|
87
|
+
self.logger.info(f"Assistant: {assistant}")
|
|
88
|
+
|
|
89
|
+
def run_query(query: str):
|
|
90
|
+
llm_response = assistant.query(query)
|
|
91
|
+
self.logger.info(f"LLM Response: {llm_response}")
|
|
92
|
+
response = RevaChatMessageResponse()
|
|
93
|
+
response.message = llm_response
|
|
94
|
+
response_queue.put(response)
|
|
95
|
+
|
|
96
|
+
t = threading.Thread(target=run_query, args=[request.message])
|
|
97
|
+
t.start()
|
|
98
|
+
|
|
99
|
+
done = False
|
|
100
|
+
while not done:
|
|
101
|
+
response = response_queue.get()
|
|
102
|
+
# We stop when we get a message and not thoughts
|
|
103
|
+
if response.message and not response.thought:
|
|
104
|
+
done = True
|
|
105
|
+
yield response
|
|
106
|
+
t.join()
|
|
107
|
+
|
|
108
|
+
def chatStream(self, request_iterator, context):
|
|
109
|
+
# TODO: This needs to be re-written to use the new callback system
|
|
110
|
+
# let's grab a reference to our assistant
|
|
111
|
+
raise NotImplementedError("chatStream is not implemented yet")
|