341ff9517a
Needed for the new fmt-check CI step to pass; the repo had never had formatting enforced before. No logic changes.
115 lines
4.2 KiB
Rust
115 lines
4.2 KiB
Rust
use anyhow::{Context, Result};
|
|
use rig_core::client::{CompletionClient, EmbeddingsClient, Nothing};
|
|
use rig_core::completion::CompletionModel;
|
|
use rig_core::embeddings::EmbeddingModel as _;
|
|
use rig_core::providers::ollama;
|
|
use serde::Deserialize;
|
|
|
|
/// Thin wrapper around a local Ollama instance via rig. Two models are used
|
|
/// deliberately: a small/cheap embedding model for the first-pass filter
|
|
/// over every article, and a larger chat model reserved for the shortlist
|
|
/// that clears the embedding threshold (see `feedsignal_core::scoring`).
|
|
pub struct Llm {
|
|
client: ollama::Client,
|
|
embedding_model: String,
|
|
embedding_dims: usize,
|
|
chat_model: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct RelevanceJudgment {
|
|
pub score: f32,
|
|
pub rationale: String,
|
|
}
|
|
|
|
impl Llm {
|
|
/// `base_url` e.g. "http://localhost:11434". Model names must already
|
|
/// be pulled locally, e.g. `ollama pull nomic-embed-text` and
|
|
/// `ollama pull llama3.1`. `embedding_dims` must match the pulled
|
|
/// embedding model (768 for nomic-embed-text, 384 for all-minilm).
|
|
pub fn new(
|
|
base_url: &str,
|
|
embedding_model: impl Into<String>,
|
|
embedding_dims: usize,
|
|
chat_model: impl Into<String>,
|
|
) -> Result<Self> {
|
|
let client = ollama::Client::builder()
|
|
.api_key(Nothing)
|
|
.base_url(base_url)
|
|
.build()
|
|
.context("failed to build ollama client")?;
|
|
Ok(Self {
|
|
client,
|
|
embedding_model: embedding_model.into(),
|
|
embedding_dims,
|
|
chat_model: chat_model.into(),
|
|
})
|
|
}
|
|
|
|
pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
|
|
let model = self
|
|
.client
|
|
.embedding_model_with_ndims(&self.embedding_model, self.embedding_dims);
|
|
let embedding = model
|
|
.embed_text(text)
|
|
.await
|
|
.context("ollama embedding request failed")?;
|
|
Ok(embedding.vec.into_iter().map(|v| v as f32).collect())
|
|
}
|
|
|
|
/// Runs the second-stage relevance judgment for one article. Called
|
|
/// only for the embedding-similarity shortlist, not every article —
|
|
/// see `feedsignal-db::shortlist_for_llm_scoring`.
|
|
pub async fn judge_relevance(
|
|
&self,
|
|
preferences_description: &str,
|
|
topic_affinities_summary: &str,
|
|
article_title: &str,
|
|
article_summary: &str,
|
|
) -> Result<RelevanceJudgment> {
|
|
let model = self.client.completion_model(&self.chat_model);
|
|
|
|
let prompt = format!(
|
|
"Preferences: {preferences_description}\n\
|
|
Historical topic affinities: {topic_affinities_summary}\n\n\
|
|
Article title: {article_title}\n\
|
|
Article summary: {article_summary}"
|
|
);
|
|
|
|
let request = model
|
|
.completion_request(prompt)
|
|
.preamble(
|
|
"You screen articles for a personal RSS reader. Given the \
|
|
reader's stated preferences, their historical topic \
|
|
affinities, and one article, respond with ONLY a JSON \
|
|
object: {\"score\": <0.0-1.0>, \"rationale\": \"<one \
|
|
sentence>\"}. score is how relevant this specific article \
|
|
is to this specific reader right now."
|
|
.to_string(),
|
|
)
|
|
.build();
|
|
|
|
let response = model
|
|
.completion(request)
|
|
.await
|
|
.context("ollama chat request failed")?;
|
|
let text: String = response
|
|
.choice
|
|
.into_iter()
|
|
.filter_map(|content| match content {
|
|
rig_core::completion::AssistantContent::Text(t) => Some(t.text),
|
|
_ => None,
|
|
})
|
|
.collect();
|
|
parse_judgment(&text)
|
|
}
|
|
}
|
|
|
|
fn parse_judgment(raw: &str) -> Result<RelevanceJudgment> {
|
|
// Models sometimes wrap JSON in prose or code fences despite instructions;
|
|
// take the outermost {...} span rather than trusting the whole response.
|
|
let start = raw.find('{').context("no JSON object in LLM response")?;
|
|
let end = raw.rfind('}').context("no JSON object in LLM response")?;
|
|
let json = &raw[start..=end];
|
|
Ok(serde_json::from_str(json)?)
|
|
}
|