llmvm_protocol/
stdio.rs

1pub use multilink::stdio::*;
2
3use multilink::{
4    jsonrpc::{JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse},
5    util::parse_from_value,
6    ProtocolError,
7};
8use serde_json::Value;
9
10use crate::service::{BackendRequest, BackendResponse, CoreRequest, CoreResponse};
11
12const GENERATION_METHOD: &str = "generation";
13const GENERATION_STREAM_METHOD: &str = "generation_stream";
14const INIT_PROJECT_METHOD: &str = "init_project";
15const GET_LAST_THREAD_INFO_METHOD: &str = "get_last_thread_info";
16const GET_ALL_THREAD_INFOS_METHOD: &str = "get_all_thread_infos";
17const GET_THREAD_MESSAGES_METHOD: &str = "get_thread_messages";
18
19impl RequestJsonRpcConvert<CoreRequest> for CoreRequest {
20    fn from_jsonrpc_request(value: JsonRpcRequest) -> Result<Option<Self>, ProtocolError> {
21        Ok(Some(match value.method.as_str() {
22            GENERATION_METHOD => CoreRequest::Generation(value.parse_params()?),
23            GENERATION_STREAM_METHOD => CoreRequest::GenerationStream(value.parse_params()?),
24            INIT_PROJECT_METHOD => CoreRequest::InitProject,
25            GET_LAST_THREAD_INFO_METHOD => CoreRequest::GetLastThreadInfo,
26            GET_ALL_THREAD_INFOS_METHOD => CoreRequest::GetAllThreadInfos,
27            GET_THREAD_MESSAGES_METHOD => CoreRequest::GetThreadMessages {
28                id: value.parse_params()?,
29            },
30            _ => return Ok(None),
31        }))
32    }
33
34    fn into_jsonrpc_request(&self) -> JsonRpcRequest {
35        let (method, params) = match self {
36            CoreRequest::Generation(request) => (
37                GENERATION_METHOD,
38                Some(serde_json::to_value(request).unwrap()),
39            ),
40            CoreRequest::GenerationStream(request) => (
41                GENERATION_STREAM_METHOD,
42                Some(serde_json::to_value(request).unwrap()),
43            ),
44            CoreRequest::InitProject => (INIT_PROJECT_METHOD, None),
45            CoreRequest::GetLastThreadInfo => (GET_LAST_THREAD_INFO_METHOD, None),
46            CoreRequest::GetAllThreadInfos => (GET_ALL_THREAD_INFOS_METHOD, None),
47            CoreRequest::GetThreadMessages { id } => (
48                GET_THREAD_MESSAGES_METHOD,
49                Some(serde_json::to_value(id).unwrap()),
50            ),
51        };
52        JsonRpcRequest::new(method.to_string(), params)
53    }
54}
55
56impl ResponseJsonRpcConvert<CoreRequest, CoreResponse> for CoreResponse {
57    fn from_jsonrpc_message(
58        value: JsonRpcMessage,
59        original_request: &CoreRequest,
60    ) -> Result<Option<Self>, ProtocolError> {
61        match value {
62            JsonRpcMessage::Response(resp) => {
63                let result = resp.get_result()?;
64                Ok(Some(match original_request {
65                    CoreRequest::Generation(_) => Self::Generation(parse_from_value(result)?),
66                    CoreRequest::GetLastThreadInfo => {
67                        Self::GetLastThreadInfo(parse_from_value(result)?)
68                    }
69                    CoreRequest::GetAllThreadInfos => {
70                        Self::GetAllThreadInfos(parse_from_value(result)?)
71                    }
72                    CoreRequest::GetThreadMessages { .. } => {
73                        Self::GetThreadMessages(parse_from_value(result)?)
74                    }
75                    CoreRequest::InitProject => Self::InitProject,
76                    _ => return Ok(None),
77                }))
78            }
79            JsonRpcMessage::Notification(resp) => {
80                let result = resp.get_result()?;
81                Ok(Some(match original_request {
82                    CoreRequest::GenerationStream(_) => {
83                        Self::GenerationStream(parse_from_value(result)?)
84                    }
85                    _ => return Ok(None),
86                }))
87            }
88            _ => Ok(None),
89        }
90    }
91
92    fn into_jsonrpc_message(response: CoreResponse, id: Value) -> JsonRpcMessage {
93        let mut is_notification = false;
94        let result = Ok(match response {
95            CoreResponse::Generation(response) => serde_json::to_value(response).unwrap(),
96            CoreResponse::GenerationStream(response) => {
97                is_notification = true;
98                serde_json::to_value(response).unwrap()
99            }
100            CoreResponse::GetLastThreadInfo(response) => serde_json::to_value(response).unwrap(),
101            CoreResponse::GetAllThreadInfos(response) => serde_json::to_value(response).unwrap(),
102            CoreResponse::GetThreadMessages(response) => serde_json::to_value(response).unwrap(),
103            CoreResponse::InitProject => Value::Null,
104        });
105        match is_notification {
106            true => JsonRpcNotification::new_with_result_params(result, id.to_string()).into(),
107            false => JsonRpcResponse::new(result, id).into(),
108        }
109    }
110}
111
112impl RequestJsonRpcConvert<BackendRequest> for BackendRequest {
113    fn from_jsonrpc_request(value: JsonRpcRequest) -> Result<Option<Self>, ProtocolError> {
114        Ok(Some(match value.method.as_str() {
115            GENERATION_METHOD => BackendRequest::Generation(value.parse_params()?),
116            GENERATION_STREAM_METHOD => BackendRequest::GenerationStream(value.parse_params()?),
117            _ => return Ok(None),
118        }))
119    }
120
121    fn into_jsonrpc_request(&self) -> JsonRpcRequest {
122        let (method, params) = match &self {
123            BackendRequest::Generation(generation_response) => (
124                GENERATION_METHOD,
125                Some(serde_json::to_value(generation_response).unwrap()),
126            ),
127            BackendRequest::GenerationStream(request) => (
128                GENERATION_STREAM_METHOD,
129                Some(serde_json::to_value(request).unwrap()),
130            ),
131        };
132        JsonRpcRequest::new(method.to_string(), params)
133    }
134}
135
136impl ResponseJsonRpcConvert<BackendRequest, BackendResponse> for BackendResponse {
137    fn from_jsonrpc_message(
138        value: JsonRpcMessage,
139        original_request: &BackendRequest,
140    ) -> Result<Option<Self>, ProtocolError> {
141        Ok(Some(match value {
142            JsonRpcMessage::Response(resp) => {
143                let result = resp.get_result()?;
144                match original_request {
145                    BackendRequest::Generation(_) => Self::Generation(parse_from_value(result)?),
146                    _ => return Ok(None),
147                }
148            }
149            JsonRpcMessage::Notification(resp) => {
150                let result = resp.get_result()?;
151                match original_request {
152                    BackendRequest::GenerationStream(_) => {
153                        Self::GenerationStream(parse_from_value(result)?)
154                    }
155                    _ => return Ok(None),
156                }
157            }
158            _ => return Ok(None),
159        }))
160    }
161
162    fn into_jsonrpc_message(response: BackendResponse, id: Value) -> JsonRpcMessage {
163        let mut is_notification = false;
164        let result = Ok(match response {
165            BackendResponse::Generation(response) => serde_json::to_value(response).unwrap(),
166            BackendResponse::GenerationStream(response) => {
167                is_notification = true;
168                serde_json::to_value(response).unwrap()
169            }
170        });
171        match is_notification {
172            true => JsonRpcNotification::new_with_result_params(result, id.to_string()).into(),
173            false => JsonRpcResponse::new(result, id).into(),
174        }
175    }
176}
llmvm_protocol/stdio.rs

llmvm_protocol/
stdio.rs
