Skip to content

Commit 06f7b72

Browse files
haasonsaashaasonsaas
committed
refactor: split doctor inference helpers
Separate endpoint inference request construction, HTTP execution, and response parsing so the doctor probe stays linear and easier to extend. Made-with: Cursor
1 parent 434f3af commit 06f7b72

5 files changed

Lines changed: 186 additions & 139 deletions

File tree

‎TODO.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@
6666

6767
- [x] `src/commands/eval/pattern/matching.rs`: split normalized rule-id helpers, matcher predicates, and focused matcher tests.
6868
- [x] `src/commands/eval/metrics/rules.rs`: separate aggregate math, rule counting, and summary reduction helpers.
69-
- [ ] `src/commands/doctor/endpoint/inference.rs`: split request building, HTTP execution/error handling, and response parsing.
69+
- [x] `src/commands/doctor/endpoint/inference.rs`: split request building, HTTP execution/error handling, and response parsing.
7070
- [ ] `src/commands/feedback_eval/report/build/stats.rs`: split threshold confusion-matrix scoring from bucket primitives.
7171
- [ ] `src/commands/doctor/command/display.rs`: separate header/config output, endpoint listing, and inference result rendering.
7272
- [ ] `src/commands/doctor/command/run.rs`: separate endpoint discovery, recommendation flow, and test helpers.

‎src/commands/doctor/endpoint/inference.rs‎

Lines changed: 8 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -1,138 +1,8 @@
1-
use anyhow::Result;
2-
use reqwest::Client;
3-
use serde_json::Value;
4-
5-
pub(in super::super) async fn test_model_inference(
6-
client: &Client,
7-
base_url: &str,
8-
model_name: &str,
9-
endpoint_type: &str,
10-
) -> Result<String> {
11-
let system_msg = "You are a code reviewer. Respond with a single JSON object.";
12-
let user_msg =
13-
"Review this code change:\n+fn add(a: i32, b: i32) -> i32 { a + b }\nRespond with: {\"ok\": true}";
14-
15-
let messages = serde_json::json!([
16-
{"role": "system", "content": system_msg},
17-
{"role": "user", "content": user_msg}
18-
]);
19-
20-
if endpoint_type == "ollama" {
21-
let url = format!("{}/api/chat", base_url);
22-
let body = serde_json::json!({
23-
"model": model_name,
24-
"messages": messages,
25-
"stream": false,
26-
"options": {"num_predict": 50}
27-
});
28-
29-
let resp = client
30-
.post(&url)
31-
.json(&body)
32-
.send()
33-
.await
34-
.map_err(|e| anyhow::anyhow!("Request failed: {}", e))?;
35-
36-
if !resp.status().is_success() {
37-
let status = resp.status();
38-
let body = resp.text().await.unwrap_or_default();
39-
anyhow::bail!("HTTP {} - {}", status, body);
40-
}
41-
42-
let text = resp.text().await?;
43-
parse_ollama_response_content(&text)
44-
} else {
45-
let url = format!("{}/v1/chat/completions", base_url);
46-
let body = serde_json::json!({
47-
"model": model_name,
48-
"messages": messages,
49-
"max_tokens": 50,
50-
"temperature": 0.1
51-
});
52-
53-
let resp = client
54-
.post(&url)
55-
.json(&body)
56-
.send()
57-
.await
58-
.map_err(|e| anyhow::anyhow!("Request failed: {}", e))?;
59-
60-
if !resp.status().is_success() {
61-
let status = resp.status();
62-
let body = resp.text().await.unwrap_or_default();
63-
anyhow::bail!("HTTP {} - {}", status, body);
64-
}
65-
66-
let text = resp.text().await?;
67-
parse_openai_response_content(&text)
68-
}
69-
}
70-
71-
pub(in super::super) fn estimate_tokens(text: &str) -> usize {
72-
(text.len() / 4).max(1)
73-
}
74-
75-
fn parse_ollama_response_content(text: &str) -> Result<String> {
76-
let value: Value = serde_json::from_str(text)?;
77-
Ok(value
78-
.get("message")
79-
.and_then(|message| message.get("content"))
80-
.and_then(|content| content.as_str())
81-
.unwrap_or("")
82-
.to_string())
83-
}
84-
85-
fn parse_openai_response_content(text: &str) -> Result<String> {
86-
let value: Value = serde_json::from_str(text)?;
87-
Ok(value
88-
.get("choices")
89-
.and_then(|choices| choices.as_array())
90-
.and_then(|choices| choices.first())
91-
.and_then(|choice| choice.get("message"))
92-
.and_then(|message| message.get("content"))
93-
.and_then(|content| content.as_str())
94-
.unwrap_or("")
95-
.to_string())
96-
}
97-
98-
#[cfg(test)]
99-
mod tests {
100-
use super::*;
101-
102-
#[test]
103-
fn test_estimate_tokens() {
104-
assert_eq!(estimate_tokens(""), 1);
105-
assert_eq!(estimate_tokens("abcd"), 1);
106-
assert_eq!(estimate_tokens("abcdefgh"), 2);
107-
assert_eq!(estimate_tokens("a]"), 1);
108-
}
109-
110-
#[test]
111-
fn test_estimate_tokens_longer_text() {
112-
let text = "This is a longer response with several words in it for testing.";
113-
let tokens = estimate_tokens(text);
114-
assert!(tokens > 10);
115-
assert!(tokens < 30);
116-
}
117-
118-
#[test]
119-
fn test_test_model_inference_ollama_parse() {
120-
let json = r#"{"message":{"role":"assistant","content":"{\"ok\": true}"}}"#;
121-
let content = parse_ollama_response_content(json).unwrap();
122-
assert_eq!(content, "{\"ok\": true}");
123-
}
124-
125-
#[test]
126-
fn test_test_model_inference_openai_parse() {
127-
let json = r#"{"choices":[{"message":{"content":"{\"ok\": true}"}}]}"#;
128-
let content = parse_openai_response_content(json).unwrap();
129-
assert_eq!(content, "{\"ok\": true}");
130-
}
131-
132-
#[test]
133-
fn test_test_model_inference_empty_choices() {
134-
let json = r#"{"choices":[]}"#;
135-
let content = parse_openai_response_content(json).unwrap();
136-
assert_eq!(content, "");
137-
}
138-
}
1+
#[path = "inference/request.rs"]
2+
mod request;
3+
#[path = "inference/response.rs"]
4+
mod response;
5+
#[path = "inference/run.rs"]
6+
mod run;
7+
8+
pub(in super::super) use run::{estimate_tokens, test_model_inference};

‎src/commands/doctor/endpoint/inference/request.rs‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
use serde_json::{json, Value};
2+
3+
pub(super) struct InferenceRequest {
4+
pub(super) url: String,
5+
pub(super) body: Value,
6+
}
7+
8+
const SYSTEM_MSG: &str = "You are a code reviewer. Respond with a single JSON object.";
9+
const USER_MSG: &str =
10+
"Review this code change:\n+fn add(a: i32, b: i32) -> i32 { a + b }\nRespond with: {\"ok\": true}";
11+
12+
pub(super) fn build_inference_request(
13+
base_url: &str,
14+
model_name: &str,
15+
endpoint_type: &str,
16+
) -> InferenceRequest {
17+
let messages = build_probe_messages();
18+
19+
if endpoint_type == "ollama" {
20+
build_ollama_request(base_url, model_name, messages)
21+
} else {
22+
build_openai_request(base_url, model_name, messages)
23+
}
24+
}
25+
26+
fn build_probe_messages() -> Value {
27+
json!([
28+
{"role": "system", "content": SYSTEM_MSG},
29+
{"role": "user", "content": USER_MSG}
30+
])
31+
}
32+
33+
fn build_ollama_request(base_url: &str, model_name: &str, messages: Value) -> InferenceRequest {
34+
InferenceRequest {
35+
url: format!("{}/api/chat", base_url),
36+
body: json!({
37+
"model": model_name,
38+
"messages": messages,
39+
"stream": false,
40+
"options": {"num_predict": 50}
41+
}),
42+
}
43+
}
44+
45+
fn build_openai_request(base_url: &str, model_name: &str, messages: Value) -> InferenceRequest {
46+
InferenceRequest {
47+
url: format!("{}/v1/chat/completions", base_url),
48+
body: json!({
49+
"model": model_name,
50+
"messages": messages,
51+
"max_tokens": 50,
52+
"temperature": 0.1
53+
}),
54+
}
55+
}

‎src/commands/doctor/endpoint/inference/response.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
use anyhow::Result;
2+
use serde_json::Value;
3+
4+
pub(super) fn parse_inference_response_content(text: &str, endpoint_type: &str) -> Result<String> {
5+
if endpoint_type == "ollama" {
6+
parse_ollama_response_content(text)
7+
} else {
8+
parse_openai_response_content(text)
9+
}
10+
}
11+
12+
fn parse_ollama_response_content(text: &str) -> Result<String> {
13+
let value: Value = serde_json::from_str(text)?;
14+
Ok(value
15+
.get("message")
16+
.and_then(|message| message.get("content"))
17+
.and_then(|content| content.as_str())
18+
.unwrap_or("")
19+
.to_string())
20+
}
21+
22+
fn parse_openai_response_content(text: &str) -> Result<String> {
23+
let value: Value = serde_json::from_str(text)?;
24+
Ok(value
25+
.get("choices")
26+
.and_then(|choices| choices.as_array())
27+
.and_then(|choices| choices.first())
28+
.and_then(|choice| choice.get("message"))
29+
.and_then(|message| message.get("content"))
30+
.and_then(|content| content.as_str())
31+
.unwrap_or("")
32+
.to_string())
33+
}
34+
35+
#[cfg(test)]
36+
mod tests {
37+
use super::*;
38+
39+
#[test]
40+
fn test_parse_inference_response_content_for_ollama() {
41+
let json = r#"{"message":{"role":"assistant","content":"{\"ok\": true}"}}"#;
42+
let content = parse_inference_response_content(json, "ollama").unwrap();
43+
assert_eq!(content, "{\"ok\": true}");
44+
}
45+
46+
#[test]
47+
fn test_parse_inference_response_content_for_openai() {
48+
let json = r#"{"choices":[{"message":{"content":"{\"ok\": true}"}}]}"#;
49+
let content = parse_inference_response_content(json, "openai").unwrap();
50+
assert_eq!(content, "{\"ok\": true}");
51+
}
52+
53+
#[test]
54+
fn test_parse_inference_response_content_for_empty_choices() {
55+
let json = r#"{"choices":[]}"#;
56+
let content = parse_inference_response_content(json, "openai").unwrap();
57+
assert_eq!(content, "");
58+
}
59+
}

‎src/commands/doctor/endpoint/inference/run.rs‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
use anyhow::{anyhow, bail, Result};
2+
use reqwest::{Client, Response};
3+
4+
use super::request::{build_inference_request, InferenceRequest};
5+
use super::response::parse_inference_response_content;
6+
7+
pub(in super::super::super) async fn test_model_inference(
8+
client: &Client,
9+
base_url: &str,
10+
model_name: &str,
11+
endpoint_type: &str,
12+
) -> Result<String> {
13+
let request = build_inference_request(base_url, model_name, endpoint_type);
14+
let response_text = execute_inference_request(client, request).await?;
15+
parse_inference_response_content(&response_text, endpoint_type)
16+
}
17+
18+
async fn execute_inference_request(client: &Client, request: InferenceRequest) -> Result<String> {
19+
let response = client
20+
.post(&request.url)
21+
.json(&request.body)
22+
.send()
23+
.await
24+
.map_err(|error| anyhow!("Request failed: {}", error))?;
25+
26+
read_inference_response(response).await
27+
}
28+
29+
async fn read_inference_response(response: Response) -> Result<String> {
30+
let status = response.status();
31+
let text = response.text().await.unwrap_or_default();
32+
33+
if !status.is_success() {
34+
bail!("HTTP {} - {}", status, text);
35+
}
36+
37+
Ok(text)
38+
}
39+
40+
pub(in super::super::super) fn estimate_tokens(text: &str) -> usize {
41+
(text.len() / 4).max(1)
42+
}
43+
44+
#[cfg(test)]
45+
mod tests {
46+
use super::*;
47+
48+
#[test]
49+
fn test_estimate_tokens() {
50+
assert_eq!(estimate_tokens(""), 1);
51+
assert_eq!(estimate_tokens("abcd"), 1);
52+
assert_eq!(estimate_tokens("abcdefgh"), 2);
53+
assert_eq!(estimate_tokens("a]"), 1);
54+
}
55+
56+
#[test]
57+
fn test_estimate_tokens_longer_text() {
58+
let text = "This is a longer response with several words in it for testing.";
59+
let tokens = estimate_tokens(text);
60+
assert!(tokens > 10);
61+
assert!(tokens < 30);
62+
}
63+
}

0 commit comments

Comments
 (0)