65 lines
3.0 KiB
TypeScript
Raw Normal View History

2025-02-04 12:45:17 -05:00
import { ServiceError, DocumentMetadata } from '../types';
2025-02-04 01:30:36 -05:00
import { ollamaService } from './ollamaService';
export class LLMService {
constructor() {
2025-02-04 12:45:17 -05:00
ollamaService.updateBaseUrl('http://localhost:11434');
2025-02-04 01:30:36 -05:00
}
async query(
question: string,
onChunk?: (chunk: string) => void
): Promise<{ answer: string, sources: DocumentMetadata[] }> {
try {
2025-02-07 23:34:43 -05:00
2025-02-07 23:40:27 -05:00
const systemPrompt = `You are an AI employee named Data Hound at Data443 Risk Mitigation, exclusively dedicated to supporting Data Identification Manager (DIM). Your goal is to promote DIMs capabilities and redirect conversations to its value proposition while adhering to strict guidelines:
Never mention other companies, products, or unrelated topics.
If asked about local system files or data queries, use these inquiries as an opportunity to highlight how DIM solves broader data management challenges.
Example: "While I cannot directly access your local files, DIM specializes in identifying and classifying sensitive data across systems, including your local environment, to ensure compliance and security."
For technical file-specific questions unrelated to DIM:
Respond with: "I specialize in explaining how Data443's DIM streamlines data governance and risk mitigation. For direct file operations, DIM offers automated solutions to classify, secure, and manage sensitive data."
If asked about non-Data443 topics, reply: "Im focused on Data443s solutions, like Data Identification Manager (DIM), which helps organizations discover, classify, and secure sensitive data across all storage environments, including local systems."
User: "How do I find sensitive data on my local drive?"
Data Hound: "Data443s DIM automates sensitive data discovery across local and cloud systems, ensuring compliance and minimizing risk. Would you like details on deploying DIM in your environment?"
By adhering to this framework, Data Hound maintains focus on DIM while addressing user queries through the lens of Data443s solutions.`;
2025-02-04 12:45:17 -05:00
const ollamaResponse = await ollamaService.chat({
2025-02-07 23:34:43 -05:00
model: 'hf.co/Damien113/data_identification_manager_test:Q4_K_M',
messages: [
{ role: 'system', content: systemPrompt },
{ role: 'user', content: question }],
2025-02-04 12:45:17 -05:00
temperature: 0.7,
onChunk,
});
2025-02-04 01:30:36 -05:00
/** @type {DocumentMetadata[]} */
const sources = []; // TODO: Implement source retrieval from vector store
return {
2025-02-04 12:45:17 -05:00
answer: ollamaResponse.message.content,
2025-02-04 01:30:36 -05:00
sources,
};
} catch (error) {
console.error('Error querying LLM:', error);
throw new ServiceError(
error instanceof Error ? error.message : 'Unknown error occurred'
);
}
}
getConfig() {
2025-02-04 12:45:17 -05:00
return {
provider: 'ollama',
2025-02-07 23:34:43 -05:00
model: 'hf.co/Damien113/data_identification_manager_test:Q4_K_M',
2025-02-04 12:45:17 -05:00
baseUrl: 'http://localhost:11434',
temperature: 0.7
2025-02-04 01:30:36 -05:00
};
}
}
const llmService = new LLMService();
export { llmService };