Skip to content

Commit a7363f4

Browse files
robdodsonrobdodson
committed
Trying different models
1 parent 0347888 commit a7363f4

File tree

1 file changed

+149
-102
lines changed

1 file changed

+149
-102
lines changed

‎src/index.ts‎

Lines changed: 149 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -26,46 +26,145 @@ interface DatabaseSchema {
2626
data_type: string;
2727
}
2828

29+
interface ColumnMetadata {
30+
description: string;
31+
examples: string[];
32+
foreignKey?: {
33+
table: string;
34+
column: string;
35+
};
36+
}
37+
2938
class OllamaMCPHost {
3039
private client: Client;
3140
private transport: StdioClientTransport;
3241
private modelName: string;
3342
private schemaCache: Map<string, DatabaseSchema[]> = new Map();
43+
private columnMetadata: Map<string, Map<string, ColumnMetadata>> = new Map();
3444
private chatHistory: { role: string; content: string }[] = [];
3545
private readonly MAX_HISTORY_LENGTH = 20;
36-
private readonly MAX_RETRIES = 2; // Maximum number of retry attempts
46+
private readonly MAX_RETRIES = 5;
47+
48+
private static readonly QUERY_GUIDELINES = `
49+
When analyzing questions:
50+
1. First write a SQL query to get the necessary information. Identify which tables contain the relevant information by looking at:
51+
- Table names and their purposes
52+
- Column names and descriptions
53+
- Foreign key relationships
54+
2. Use the 'query' tool to execute the SQL query
55+
3. If unsure about table contents, write a sample query first:
56+
SELECT column_name, COUNT(*) FROM table_name GROUP BY column_name LIMIT 5;
57+
4. For complex questions, break down into multiple queries:
58+
- First query to validate data availability
59+
- Second query to get detailed information
60+
5. Always include appropriate JOIN conditions when combining tables
61+
6. Use WHERE clauses to filter irrelevant data
62+
7. Consider using ORDER BY for sorted results
63+
64+
Important: Only use SELECT statements - no modifications allowed!
65+
66+
When you are finished, analyze the results and provide a natural language response.`;
3767

3868
constructor(modelName?: string) {
3969
this.modelName =
4070
modelName || process.env.OLLAMA_MODEL || "qwen2.5-coder:7b-instruct";
41-
4271
this.transport = new StdioClientTransport({
4372
command: "npx",
4473
args: ["-y", "@modelcontextprotocol/server-postgres", databaseUrl!],
4574
});
46-
4775
this.client = new Client(
48-
{
49-
name: "ollama-mcp-host",
50-
version: "1.0.0",
51-
},
52-
{
53-
capabilities: {},
54-
}
76+
{ name: "ollama-mcp-host", version: "1.0.0" },
77+
{ capabilities: {} }
5578
);
5679
}
5780

58-
private addToHistory(role: string, content: string) {
59-
this.chatHistory.push({ role, content });
81+
private async detectTableRelationships(): Promise<void> {
82+
// Query the database to find foreign key relationships
83+
const sql = `
84+
SELECT
85+
tc.table_name as table_name,
86+
kcu.column_name as column_name,
87+
ccu.table_name AS foreign_table_name,
88+
ccu.column_name AS foreign_column_name
89+
FROM information_schema.table_constraints tc
90+
JOIN information_schema.key_column_usage kcu
91+
ON tc.constraint_name = kcu.constraint_name
92+
JOIN information_schema.constraint_column_usage ccu
93+
ON ccu.constraint_name = tc.constraint_name
94+
WHERE constraint_type = 'FOREIGN KEY'
95+
`;
6096

61-
while (this.chatHistory.length > this.MAX_HISTORY_LENGTH) {
62-
this.chatHistory.splice(0, 2);
97+
try {
98+
const result = await this.executeQuery(sql);
99+
const relationships = JSON.parse(result);
100+
101+
// Create initial metadata for foreign keys
102+
relationships.forEach((rel: any) => {
103+
const tableMetadata =
104+
this.columnMetadata.get(rel.table_name) || new Map();
105+
106+
tableMetadata.set(rel.column_name, {
107+
description: `Foreign key referencing ${rel.foreign_table_name}.${rel.foreign_column_name}`,
108+
examples: [],
109+
foreignKey: {
110+
table: rel.foreign_table_name,
111+
column: rel.foreign_column_name,
112+
},
113+
});
114+
115+
this.columnMetadata.set(rel.table_name, tableMetadata);
116+
});
117+
} catch (error) {
118+
console.error("Error detecting table relationships:", error);
119+
}
120+
}
121+
122+
private buildSystemPrompt(includeErrorContext: string = ""): string {
123+
let prompt =
124+
"You are a data analyst assistant. You have access to a PostgreSQL database with these tables:\n\n";
125+
126+
// Add detailed schema information
127+
for (const [tableName, schema] of this.schemaCache.entries()) {
128+
prompt += `Table: ${tableName}\n`;
129+
prompt += "Columns:\n";
130+
131+
for (const column of schema) {
132+
const metadata = this.columnMetadata
133+
.get(tableName)
134+
?.get(column.column_name);
135+
prompt += `- ${column.column_name} (${column.data_type})`;
136+
137+
if (metadata) {
138+
prompt += `: ${metadata.description}`;
139+
if (metadata.foreignKey) {
140+
prompt += ` [References ${metadata.foreignKey.table}.${metadata.foreignKey.column}]`;
141+
}
142+
}
143+
prompt += "\n";
144+
}
145+
prompt += "\n";
146+
}
147+
148+
// Add query guidelines
149+
prompt += "\nQuery Guidelines:\n";
150+
prompt += OllamaMCPHost.QUERY_GUIDELINES;
151+
152+
if (includeErrorContext) {
153+
prompt += `\nPrevious Error Context: ${includeErrorContext}\n`;
154+
prompt +=
155+
"Please revise your approach and try a different query strategy.\n";
63156
}
157+
158+
return prompt;
64159
}
65160

66161
async connect() {
67162
await this.client.connect(this.transport);
68163

164+
// First detect relationships
165+
await this.detectTableRelationships();
166+
167+
// Then load schemas
69168
const resources = await this.client.request(
70169
{ method: "resources/list" },
71170
ListResourcesResultSchema
@@ -94,42 +193,11 @@ class OllamaMCPHost {
94193
error instanceof Error ? error.message : String(error)
95194
);
96195
}
97-
} else {
98-
console.warn(`No text content found for resource ${resource.uri}`);
99196
}
100197
}
101198
}
102199
}
103200

104-
private buildSystemPrompt(includeErrorContext: string = ""): string {
105-
let prompt =
106-
"You are a data analyst assistant. You have access to a PostgreSQL database with the following tables and schemas:\n\n";
107-
108-
for (const [tableName, schema] of this.schemaCache.entries()) {
109-
prompt += `Table: ${tableName}\nColumns:\n`;
110-
for (const column of schema) {
111-
prompt += `- ${column.column_name} (${column.data_type})\n`;
112-
}
113-
prompt += "\n";
114-
}
115-
116-
prompt += "\nWhen answering questions about the data:\n";
117-
prompt += "1. First write a SQL query to get the necessary information\n";
118-
prompt += "2. Use the 'query' tool to execute the SQL query\n";
119-
prompt +=
120-
"3. Analyze the results and provide a natural language response\n";
121-
prompt +=
122-
"\nImportant: Only use SELECT statements - no modifications allowed.\n";
123-
124-
if (includeErrorContext) {
125-
prompt += `\nThe previous query attempt failed with the following error: ${includeErrorContext}\n`;
126-
prompt +=
127-
"Please revise your approach and try a different query strategy.\n";
128-
}
129-
130-
return prompt;
131-
}
132-
133201
private async executeQuery(sql: string): Promise<string> {
134202
const response = await this.client.request(
135203
{
@@ -148,47 +216,10 @@ class OllamaMCPHost {
148216
return response.content[0].text as string;
149217
}
150218

151-
private async attemptQuery(
152-
messages: { role: string; content: string }[]
153-
): Promise<{
154-
success: boolean;
155-
response: string;
156-
sql?: string;
157-
queryResult?: string;
158-
error?: string;
159-
}> {
160-
const response = await ollama.chat({
161-
model: this.modelName,
162-
messages: messages,
163-
});
164-
165-
const sqlMatch = response.message.content.match(/```sql\n([\s\S]*?)\n```/);
166-
if (!sqlMatch) {
167-
return {
168-
success: false,
169-
response: response.message.content,
170-
error: "No SQL query found in response",
171-
};
172-
}
173-
174-
const sql = sqlMatch[1].trim();
175-
console.log("Executing SQL:", sql);
176-
177-
try {
178-
const queryResult = await this.executeQuery(sql);
179-
return {
180-
success: true,
181-
response: response.message.content,
182-
sql,
183-
queryResult,
184-
};
185-
} catch (error) {
186-
return {
187-
success: false,
188-
response: response.message.content,
189-
sql,
190-
error: error instanceof Error ? error.message : String(error),
191-
};
219+
private addToHistory(role: string, content: string) {
220+
this.chatHistory.push({ role, content });
221+
while (this.chatHistory.length > this.MAX_HISTORY_LENGTH) {
222+
this.chatHistory.shift();
192223
}
193224
}
194225

@@ -198,7 +229,6 @@ class OllamaMCPHost {
198229
let lastError: string | undefined;
199230

200231
while (attemptCount <= this.MAX_RETRIES) {
201-
// Prepare messages for this attempt
202232
const messages = [
203233
{ role: "system", content: this.buildSystemPrompt(lastError) },
204234
...this.chatHistory,
@@ -213,30 +243,47 @@ class OllamaMCPHost {
213243
attemptCount > 0 ? `\nRetry attempt ${attemptCount}...` : ""
214244
);
215245

216-
const result = await this.attemptQuery(messages);
246+
// Get response from Ollama
247+
const response = await ollama.chat({
248+
model: this.modelName,
249+
messages: messages,
250+
});
251+
252+
// Extract SQL query
253+
const sqlMatch = response.message.content.match(
254+
/```sql\n([\s\S]*?)\n```/
255+
);
256+
if (!sqlMatch) {
257+
return response.message.content;
258+
}
259+
260+
const sql = sqlMatch[1].trim();
261+
console.log("Executing SQL:", sql);
217262

218-
if (result.success && result.queryResult) {
219-
// Query succeeded
220-
this.addToHistory("assistant", result.response);
263+
try {
264+
// Execute the query
265+
const queryResult = await this.executeQuery(sql);
266+
this.addToHistory("assistant", response.message.content);
221267

222-
// Request interpretation of results
223-
const resultPrompt = `Here are the results of the SQL query: ${result.queryResult}\n\nPlease analyze these results and provide a clear summary.`;
224-
this.addToHistory("user", resultPrompt);
268+
// Ask for result interpretation
269+
const interpretationMessages = [
270+
...messages,
271+
{ role: "assistant", content: response.message.content },
272+
{
273+
role: "user",
274+
content: `Here are the results of the SQL query: ${queryResult}\n\nPlease analyze these results and provide a clear summary.`,
275+
},
276+
];
225277

226278
const finalResponse = await ollama.chat({
227279
model: this.modelName,
228-
messages: [
229-
...messages,
230-
{ role: "assistant", content: result.response },
231-
{ role: "user", content: resultPrompt },
232-
],
280+
messages: interpretationMessages,
233281
});
234282

235283
this.addToHistory("assistant", finalResponse.message.content);
236284
return finalResponse.message.content;
237-
} else {
238-
// Query failed
239-
lastError = result.error;
285+
} catch (error) {
286+
lastError = error instanceof Error ? error.message : String(error);
240287
if (attemptCount === this.MAX_RETRIES) {
241288
return `I apologize, but I was unable to successfully query the database after ${
242289
this.MAX_RETRIES + 1

0 commit comments

Comments
 (0)