Révisions pour rag_hybrid.py — Búsqueda híbrida FTS5 + embeddings

1

+

#!/usr/bin/env python3

2

+

"""

3

+

rag_hybrid.py — Búsqueda híbrida (FTS5 + cosine similarity) para la base de conocimiento RAG.

4

+

5

+

Combina lo mejor de dos mundos:

6

+

- FTS5: búsqueda textual exacta con stemming y diacríticos

7

+

- Embeddings: búsqueda semántica por cosine similarity

8

+

9

+

La fórmula de combinación es:

10

+

score_total = alpha * sigmoid(fts_score/10) + (1 - alpha) * cosine_sim

11

+

12

+

Uso:

13

+

python3 rag_hybrid.py "cómo configurar un proxy inverso"

14

+

python3 rag_hybrid.py "hooks de git" --alpha 0.6 --limite 10

15

+

"""

16

+

17

+

import argparse

18

+

import math

19

+

import os

20

+

import sqlite3

21

+

import sys

22

+

from typing import Any

23

+

24

+

import numpy as np

25

+

from numpy.linalg import norm

26

+

27

+

DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")

28

+

EMBEDDING_DIM = 1024

29

+

EMBEDDING_BYTES = EMBEDDING_DIM * 4 # 4096 bytes

30

+

31

+

32

+

def blob_to_vector(blob: bytes) -> np.ndarray:

33

+

if len(blob) != EMBEDDING_BYTES:

34

+

raise ValueError(f"BLOB incorrecto: {len(blob)} bytes (esperados {EMBEDDING_BYTES})")

35

+

return np.frombuffer(blob, dtype=np.float32).copy()

36

+

37

+

38

+

def normalizar(vector: np.ndarray) -> np.ndarray:

39

+

norma = norm(vector)

40

+

return vector / norma if norma > 0 else vector

41

+

42

+

43

+

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:

44

+

return float(np.dot(a, b))

45

+

46

+

47

+

def sigmoid_fts(fts_score: float) -> float:

48

+

return 1.0 / (1.0 + math.exp(-fts_score / 10.0))

49

+

50

+

51

+

def obtener_embedding_ollama(texto: str) -> np.ndarray:

52

+

import requests

53

+

response = requests.post(

54

+

"http://localhost:11434/api/embeddings",

55

+

json={"model": "bge-m3", "prompt": texto},

56

+

timeout=30,

57

+

)

58

+

response.raise_for_status()

59

+

return np.array(response.json()["embedding"], dtype=np.float32)

60

+

61

+

62

+

def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 10, tag: str | None = None) -> list[dict]:

63

+

query_emb = obtener_embedding_ollama(query)

64

+

query_emb = normalizar(query_emb)

65

+

66

+

if not os.path.exists(DB_PATH):

67

+

print(f"Error: BD no encontrada en {DB_PATH}", file=sys.stderr)

68

+

return []

69

+

70

+

conn = sqlite3.connect(DB_PATH)

71

+

conn.row_factory = sqlite3.Row

72

+

73

+

where_clauses: list[str] = ["chunks_fts MATCH ?"]

74

+

params: list[Any] = [query]

75

+

76

+

if tag:

77

+

where_clauses.append("c.tags LIKE ?")

78

+

params.append(f"%{tag}%")

79

+

80

+

sql = f"""

81

+

SELECT c.id, c.content, e.vector, c.doc_id,

82

+

bm25(chunks_fts) as fts_score

83

+

FROM chunks_fts

84

+

JOIN chunks c ON chunks_fts.rowid = c.id

85

+

JOIN embeddings e ON e.chunk_id = c.id

86

+

WHERE {" AND ".join(where_clauses)}

87

+

ORDER BY fts_score DESC

88

+

LIMIT ?

89

+

"""

90

+

params.append(top_k * 2)

91

+

92

+

cursor = conn.execute(sql, params)

93

+

rows = cursor.fetchall()

94

+

conn.close()

95

+

96

+

if not rows:

97

+

print(f"Sin resultados para: {query}")

98

+

return []

99

+

100

+

results = []

101

+

for row in rows:

102

+

try:

103

+

emb_vec = blob_to_vector(row["vector"])

104

+

except ValueError as e:

105

+

print(f"[aviso] {e}", file=sys.stderr)

106

+

continue

107

+

108

+

emb_vec = normalizar(emb_vec)

109

+

cos_sim = cosine_similarity(query_emb, emb_vec)

110

+

fts_norm = sigmoid_fts(row["fts_score"])

111

+

score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim

112

+

113

+

results.append({

114

+

"contenido": row["content"],

115

+

"doc_id": row["doc_id"],

116

+

"score": score_total,

117

+

"fts_score": round(row["fts_score"], 4),

118

+

"cos_sim": round(cos_sim, 4),

119

+

})

120

+

121

+

results.sort(key=lambda x: x["score"], reverse=True)

122

+

return results[:top_k]

123

+

124

+

125

+

def main():

126

+

parser = argparse.ArgumentParser(description="Búsqueda híbrida FTS5 + embeddings")

127

+

parser.add_argument("consulta", help="Texto de la consulta")

128

+

parser.add_argument("--alpha", type=float, default=0.4, help="Peso de FTS5 (0-1)")

129

+

parser.add_argument("--limite", type=int, default=10, help="Número de resultados")

130

+

parser.add_argument("--tag", help="Filtrar por etiqueta")

131

+

args = parser.parse_args()

132

+

133

+

resultados = hybrid_search(query=args.consulta, alpha=args.alpha, top_k=args.limite, tag=args.tag)

134

+

if not resultados:

135

+

sys.exit(0)

136

+

137

+

print(f"\n{= * 70}")

138

+

print(f" Consulta: {args.consulta} | Alpha: {args.alpha} | Resultados: {len(resultados)}")

139

+

print(f"{= * 70}\n")

140

+

141

+

for i, r in enumerate(resultados, 1):

142

+

preview = r["contenido"].replace("\n", " ")[:120]

143

+

print(f" [{i:2d}] Score: {r[score]:.4f} (FTS: {r[fts_score]:.4f} | Cos: {r[cos_sim]:.4f})")

144

+

print(f" {preview}\n")

145

+

146

+

print(f"{= * 70}")

147

+

148

+

149

+

if __name__ == "__main__":

150

+

main()

atareao / rag_hybrid.py — Búsqueda híbrida FTS5 + embeddings

atareao a révisé ce gist 3 weeks ago. Aller à la révision

		@@ -0,0 +1,150 @@
1	+	#!/usr/bin/env python3
2	+	"""
3	+	rag_hybrid.py — Búsqueda híbrida (FTS5 + cosine similarity) para la base de conocimiento RAG.
4	+
5	+	Combina lo mejor de dos mundos:
6	+	- FTS5: búsqueda textual exacta con stemming y diacríticos
7	+	- Embeddings: búsqueda semántica por cosine similarity
8	+
9	+	La fórmula de combinación es:
10	+	score_total = alpha * sigmoid(fts_score/10) + (1 - alpha) * cosine_sim
11	+
12	+	Uso:
13	+	python3 rag_hybrid.py "cómo configurar un proxy inverso"
14	+	python3 rag_hybrid.py "hooks de git" --alpha 0.6 --limite 10
15	+	"""
16	+
17	+	import argparse
18	+	import math
19	+	import os
20	+	import sqlite3
21	+	import sys
22	+	from typing import Any
23	+
24	+	import numpy as np
25	+	from numpy.linalg import norm
26	+
27	+	DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")
28	+	EMBEDDING_DIM = 1024
29	+	EMBEDDING_BYTES = EMBEDDING_DIM * 4 # 4096 bytes
30	+
31	+
32	+	def blob_to_vector(blob: bytes) -> np.ndarray:
33	+	if len(blob) != EMBEDDING_BYTES:
34	+	raise ValueError(f"BLOB incorrecto: {len(blob)} bytes (esperados {EMBEDDING_BYTES})")
35	+	return np.frombuffer(blob, dtype=np.float32).copy()
36	+
37	+
38	+	def normalizar(vector: np.ndarray) -> np.ndarray:
39	+	norma = norm(vector)
40	+	return vector / norma if norma > 0 else vector
41	+
42	+
43	+	def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
44	+	return float(np.dot(a, b))
45	+
46	+
47	+	def sigmoid_fts(fts_score: float) -> float:
48	+	return 1.0 / (1.0 + math.exp(-fts_score / 10.0))
49	+
50	+
51	+	def obtener_embedding_ollama(texto: str) -> np.ndarray:
52	+	import requests
53	+	response = requests.post(
54	+	"http://localhost:11434/api/embeddings",
55	+	json={"model": "bge-m3", "prompt": texto},
56	+	timeout=30,
57	+	)
58	+	response.raise_for_status()
59	+	return np.array(response.json()["embedding"], dtype=np.float32)
60	+
61	+
62	+	def hybrid_search(query: str, alpha: float = 0.4, top_k: int = 10, tag: str \| None = None) -> list[dict]:
63	+	query_emb = obtener_embedding_ollama(query)
64	+	query_emb = normalizar(query_emb)
65	+
66	+	if not os.path.exists(DB_PATH):
67	+	print(f"Error: BD no encontrada en {DB_PATH}", file=sys.stderr)
68	+	return []
69	+
70	+	conn = sqlite3.connect(DB_PATH)
71	+	conn.row_factory = sqlite3.Row
72	+
73	+	where_clauses: list[str] = ["chunks_fts MATCH ?"]
74	+	params: list[Any] = [query]
75	+
76	+	if tag:
77	+	where_clauses.append("c.tags LIKE ?")
78	+	params.append(f"%{tag}%")
79	+
80	+	sql = f"""
81	+	SELECT c.id, c.content, e.vector, c.doc_id,
82	+	bm25(chunks_fts) as fts_score
83	+	FROM chunks_fts
84	+	JOIN chunks c ON chunks_fts.rowid = c.id
85	+	JOIN embeddings e ON e.chunk_id = c.id
86	+	WHERE {" AND ".join(where_clauses)}
87	+	ORDER BY fts_score DESC
88	+	LIMIT ?
89	+	"""
90	+	params.append(top_k * 2)
91	+
92	+	cursor = conn.execute(sql, params)
93	+	rows = cursor.fetchall()
94	+	conn.close()
95	+
96	+	if not rows:
97	+	print(f"Sin resultados para: {query}")
98	+	return []
99	+
100	+	results = []
101	+	for row in rows:
102	+	try:
103	+	emb_vec = blob_to_vector(row["vector"])
104	+	except ValueError as e:
105	+	print(f"[aviso] {e}", file=sys.stderr)
106	+	continue
107	+
108	+	emb_vec = normalizar(emb_vec)
109	+	cos_sim = cosine_similarity(query_emb, emb_vec)
110	+	fts_norm = sigmoid_fts(row["fts_score"])
111	+	score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim
112	+
113	+	results.append({
114	+	"contenido": row["content"],
115	+	"doc_id": row["doc_id"],
116	+	"score": score_total,
117	+	"fts_score": round(row["fts_score"], 4),
118	+	"cos_sim": round(cos_sim, 4),
119	+	})
120	+
121	+	results.sort(key=lambda x: x["score"], reverse=True)
122	+	return results[:top_k]
123	+
124	+
125	+	def main():
126	+	parser = argparse.ArgumentParser(description="Búsqueda híbrida FTS5 + embeddings")
127	+	parser.add_argument("consulta", help="Texto de la consulta")
128	+	parser.add_argument("--alpha", type=float, default=0.4, help="Peso de FTS5 (0-1)")
129	+	parser.add_argument("--limite", type=int, default=10, help="Número de resultados")
130	+	parser.add_argument("--tag", help="Filtrar por etiqueta")
131	+	args = parser.parse_args()
132	+
133	+	resultados = hybrid_search(query=args.consulta, alpha=args.alpha, top_k=args.limite, tag=args.tag)
134	+	if not resultados:
135	+	sys.exit(0)
136	+
137	+	print(f"\n{= * 70}")
138	+	print(f" Consulta: {args.consulta} \| Alpha: {args.alpha} \| Resultados: {len(resultados)}")
139	+	print(f"{= * 70}\n")
140	+
141	+	for i, r in enumerate(resultados, 1):
142	+	preview = r["contenido"].replace("\n", " ")[:120]
143	+	print(f" [{i:2d}] Score: {r[score]:.4f} (FTS: {r[fts_score]:.4f} \| Cos: {r[cos_sim]:.4f})")
144	+	print(f" {preview}\n")
145	+
146	+	print(f"{= * 70}")
147	+
148	+
149	+	if __name__ == "__main__":
150	+	main()