Merge branch 'feature/ai-translator' into develop
This commit is contained in:
@@ -7,7 +7,7 @@ import (
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
"tm/internal/tender"
|
||||
"tm/internal/notice"
|
||||
"tm/pkg/config"
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
@@ -78,7 +78,7 @@ func InitMongoDB(conf config.MongoConfig, log logger.Logger) *mongo.ConnectionMa
|
||||
// init TED scraper
|
||||
func InitTEDScraper(config Config, mongoManager *mongo.ConnectionManager, appLogger logger.Logger, notify notification.SDK) {
|
||||
// Initialize tender repository
|
||||
tenderRepo := tender.NewRepository(mongoManager, appLogger)
|
||||
noticeRepo := notice.NewRepository(mongoManager, appLogger)
|
||||
|
||||
// Initialize TED scraper
|
||||
tedScraper := ted.NewTEDScraper(
|
||||
@@ -92,15 +92,16 @@ func InitTEDScraper(config Config, mongoManager *mongo.ConnectionManager, appLog
|
||||
DownloadDir: config.TED.DownloadDir,
|
||||
CleanupAfter: config.TED.CleanupAfter,
|
||||
ScrapingInterval: config.TED.ScrapingInterval,
|
||||
AlertMail: config.AlertMail,
|
||||
},
|
||||
appLogger,
|
||||
mongoManager,
|
||||
tenderRepo,
|
||||
noticeRepo,
|
||||
notify,
|
||||
)
|
||||
|
||||
// start TED scraper job
|
||||
schedule.NewCronScheduler(appLogger, mongoManager, tenderRepo).AddJob(schedule.Job{
|
||||
schedule.NewCronScheduler(appLogger, mongoManager, noticeRepo).AddJob(schedule.Job{
|
||||
Name: "TED Scraper Job",
|
||||
Func: func() {
|
||||
_ = tedScraper.Run(context.Background())
|
||||
|
||||
@@ -11,6 +11,7 @@ type Config struct {
|
||||
Logging config.LoggingConfig
|
||||
TED ScraperConfig
|
||||
Notification config.NotificationConfig
|
||||
AlertMail string `env:"ALERT_MAIL" envDefault:"alerts@opplens.com"`
|
||||
}
|
||||
|
||||
type ScraperConfig struct {
|
||||
|
||||
+3
-5
@@ -22,8 +22,6 @@ func main() {
|
||||
"version": "1.0.0",
|
||||
})
|
||||
|
||||
// Initialize notification service
|
||||
|
||||
// Initialize MongoDB connection
|
||||
mongoManager := bootstrap.InitMongoDB(config.Database.MongoDB, appLogger)
|
||||
defer func() {
|
||||
@@ -36,7 +34,10 @@ func main() {
|
||||
}
|
||||
}()
|
||||
|
||||
// Initialize notification service
|
||||
notificationService := bootstrap.InitNotificationService(config.Notification, appLogger)
|
||||
|
||||
// Initialize TED scraper
|
||||
bootstrap.InitTEDScraper(*config, mongoManager, appLogger, notificationService)
|
||||
|
||||
// Set up signal handling for graceful shutdown
|
||||
@@ -47,8 +48,5 @@ func main() {
|
||||
<-signalChan
|
||||
appLogger.Info("Received shutdown signal, stopping scheduler...", map[string]interface{}{})
|
||||
|
||||
// Stop the scheduler
|
||||
// scheduler.Stop()
|
||||
|
||||
appLogger.Info("TED scraper application stopped", map[string]interface{}{})
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
"tm/pkg/notification"
|
||||
"tm/pkg/ollama"
|
||||
"tm/pkg/redis"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
@@ -213,3 +214,38 @@ func InitNotificationService(conf config.NotificationConfig, log logger.Logger)
|
||||
|
||||
return *notificationSDK
|
||||
}
|
||||
|
||||
// Init Ollama Service
|
||||
func InitOllamaService(conf config.OllamaConfig, log logger.Logger) ollama.SDK {
|
||||
ollamaConfig := &ollama.Config{
|
||||
BaseURL: conf.BaseURL,
|
||||
Timeout: conf.Timeout,
|
||||
RetryAttempts: conf.RetryAttempts,
|
||||
RetryDelay: conf.RetryDelay,
|
||||
EnableLogging: conf.EnableLogging,
|
||||
UserAgent: conf.UserAgent,
|
||||
DefaultModel: conf.DefaultModel,
|
||||
DefaultTemperature: conf.DefaultTemperature,
|
||||
DefaultTopP: conf.DefaultTopP,
|
||||
DefaultMaxTokens: conf.DefaultMaxTokens,
|
||||
EnableStreaming: conf.EnableStreaming,
|
||||
}
|
||||
|
||||
ollamaSDK, err := ollama.New(ollamaConfig, log)
|
||||
if err != nil {
|
||||
log.Error("Failed to initialize Ollama SDK", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
log.Info("Ollama SDK initialized successfully", map[string]interface{}{
|
||||
"base_url": conf.BaseURL,
|
||||
"timeout": conf.Timeout,
|
||||
"retry_attempts": conf.RetryAttempts,
|
||||
"retry_delay": conf.RetryDelay,
|
||||
"enable_logging": conf.EnableLogging,
|
||||
"user_agent": conf.UserAgent,
|
||||
})
|
||||
|
||||
return *ollamaSDK
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ type Config struct {
|
||||
RateLimit config.RateLimitConfig
|
||||
Assets config.AssetsConfig
|
||||
Notification config.NotificationConfig
|
||||
Ollama config.OllamaConfig
|
||||
UserAuth AuthConfig
|
||||
CustomerAuth AuthConfig
|
||||
}
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
package bootstrap
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
"tm/cmd/worker/workers"
|
||||
"tm/internal/notice"
|
||||
"tm/internal/tender"
|
||||
"tm/pkg/config"
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
"tm/pkg/notification"
|
||||
"tm/pkg/ollama"
|
||||
"tm/pkg/schedule"
|
||||
)
|
||||
|
||||
// Init Application Configuration
|
||||
func InitConfig() (*Config, error) {
|
||||
conf, err := config.LoadConfig(".", &Config{})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load config: %v", err)
|
||||
}
|
||||
|
||||
return conf, nil
|
||||
}
|
||||
|
||||
// Init Logger Service
|
||||
func InitLogger(conf config.LoggingConfig) logger.Logger {
|
||||
return logger.NewLogger(&logger.Config{
|
||||
Level: conf.Level,
|
||||
Format: conf.Format,
|
||||
Output: conf.Output,
|
||||
File: logger.FileConfig{
|
||||
Path: conf.File.Path,
|
||||
MaxSize: conf.File.MaxSize,
|
||||
MaxBackups: conf.File.MaxBackups,
|
||||
MaxAge: conf.File.MaxAge,
|
||||
Compress: conf.File.Compress,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Init MongoDB Connection Manager
|
||||
func InitMongoDB(conf config.MongoConfig, log logger.Logger) *mongo.ConnectionManager {
|
||||
// Convert infra.MongoConfig to mongo.ConnectionConfig
|
||||
connectionConfig := mongo.ConnectionConfig{
|
||||
URI: conf.URI,
|
||||
Database: conf.Name,
|
||||
MaxPoolSize: uint64(conf.MaxPoolSize),
|
||||
MinPoolSize: 5, // Default minimum pool size
|
||||
MaxConnIdleTime: 30 * time.Minute, // Default idle time
|
||||
ConnectTimeout: conf.Timeout,
|
||||
ServerSelectionTimeout: 5 * time.Second, // Default server selection timeout
|
||||
}
|
||||
|
||||
// Create MongoDB connection manager
|
||||
connectionManager, err := mongo.NewConnectionManager(connectionConfig, log)
|
||||
if err != nil {
|
||||
log.Error("Failed to initialize MongoDB connection", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
"uri": conf.URI,
|
||||
"db": conf.Name,
|
||||
})
|
||||
panic(fmt.Sprintf("Failed to initialize MongoDB connection: %v", err))
|
||||
}
|
||||
|
||||
log.Info("MongoDB connection manager initialized successfully", map[string]interface{}{
|
||||
"database": conf.Name,
|
||||
"uri": conf.URI,
|
||||
"pool_size": conf.MaxPoolSize,
|
||||
})
|
||||
|
||||
return connectionManager
|
||||
}
|
||||
|
||||
// Init Worker
|
||||
func InitWorker(config Config, mongoManager *mongo.ConnectionManager, appLogger logger.Logger, notify notification.SDK, ollamaSDK *ollama.SDK) {
|
||||
// Initialize repositories
|
||||
noticeRepo := notice.NewRepository(mongoManager, appLogger)
|
||||
tenderRepo := tender.NewRepository(mongoManager, appLogger)
|
||||
|
||||
worker := workers.NewNoticeWorker(mongoManager, appLogger, notify, noticeRepo, tenderRepo, ollamaSDK)
|
||||
worker.Run()
|
||||
|
||||
// start Worker job
|
||||
schedule.NewCronScheduler(appLogger, mongoManager, noticeRepo).AddJob(schedule.Job{
|
||||
Name: "Worker Job",
|
||||
Func: func() {
|
||||
worker := workers.NewNoticeWorker(mongoManager, appLogger, notify, noticeRepo, tenderRepo, ollamaSDK)
|
||||
worker.Run()
|
||||
},
|
||||
Expr: config.Worker.Interval,
|
||||
})
|
||||
|
||||
}
|
||||
|
||||
// Init Notification Service
|
||||
func InitNotificationService(conf config.NotificationConfig, log logger.Logger) notification.SDK {
|
||||
notificationConfig := ¬ification.Config{
|
||||
BaseURL: conf.BaseURL,
|
||||
Timeout: conf.Timeout,
|
||||
RetryAttempts: conf.RetryAttempts,
|
||||
RetryDelay: conf.RetryDelay,
|
||||
EnableLogging: conf.EnableLogging,
|
||||
UserAgent: conf.UserAgent,
|
||||
}
|
||||
|
||||
notificationSDK, err := notification.New(notificationConfig, log)
|
||||
if err != nil {
|
||||
log.Error("Failed to initialize Notification SDK", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
fmt.Printf("Failed to initialize Notification SDK: %v\n", err.Error())
|
||||
}
|
||||
|
||||
log.Info("Notification SDK initialized successfully", map[string]interface{}{
|
||||
"base_url": conf.BaseURL,
|
||||
"timeout": conf.Timeout,
|
||||
"retry_attempts": conf.RetryAttempts,
|
||||
"retry_delay": conf.RetryDelay,
|
||||
"enable_logging": conf.EnableLogging,
|
||||
"user_agent": conf.UserAgent,
|
||||
})
|
||||
|
||||
return *notificationSDK
|
||||
}
|
||||
|
||||
// Init Ollama Service
|
||||
func InitOllamaService(conf config.OllamaConfig, log logger.Logger) *ollama.SDK {
|
||||
ollamaSDK, err := ollama.New(&ollama.Config{
|
||||
BaseURL: conf.BaseURL,
|
||||
Timeout: conf.Timeout,
|
||||
RetryAttempts: conf.RetryAttempts,
|
||||
RetryDelay: conf.RetryDelay,
|
||||
EnableLogging: conf.EnableLogging,
|
||||
UserAgent: conf.UserAgent,
|
||||
DefaultModel: conf.DefaultModel,
|
||||
DefaultTemperature: conf.DefaultTemperature,
|
||||
DefaultTopP: conf.DefaultTopP,
|
||||
DefaultMaxTokens: conf.DefaultMaxTokens,
|
||||
EnableStreaming: conf.EnableStreaming,
|
||||
}, log)
|
||||
if err != nil {
|
||||
log.Error("Failed to initialize Ollama SDK", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
return ollamaSDK
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package bootstrap
|
||||
|
||||
import (
|
||||
"tm/pkg/config"
|
||||
)
|
||||
|
||||
// Config defines the scraper application configuration
|
||||
type Config struct {
|
||||
Database config.DatabaseConfig
|
||||
Logging config.LoggingConfig
|
||||
Notification config.NotificationConfig
|
||||
Ollama config.OllamaConfig
|
||||
Worker WorkerConfig
|
||||
AlertMail string `env:"ALERT_MAIL" envDefault:"alerts@opplens.com"`
|
||||
}
|
||||
|
||||
type WorkerConfig struct {
|
||||
Interval string `env:"WORKER_INTERVAL" envDefault:"* 10 * * * *"`
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
"tm/cmd/worker/bootstrap"
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Load configuration
|
||||
config, err := bootstrap.InitConfig()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Initialize logger
|
||||
appLogger := bootstrap.InitLogger(config.Logging)
|
||||
appLogger.Info("Starting worker application", map[string]interface{}{
|
||||
"version": "1.0.0",
|
||||
})
|
||||
|
||||
// Initialize MongoDB connection
|
||||
mongoManager := bootstrap.InitMongoDB(config.Database.MongoDB, appLogger)
|
||||
defer func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := mongoManager.Close(ctx); err != nil {
|
||||
appLogger.Error("Failed to close MongoDB connection", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}()
|
||||
|
||||
// Initialize notification service
|
||||
notificationService := bootstrap.InitNotificationService(config.Notification, appLogger)
|
||||
|
||||
// Initialize ollama service
|
||||
ollamaSDK := bootstrap.InitOllamaService(config.Ollama, appLogger)
|
||||
aiModels, err := ollamaSDK.ListModels(context.Background())
|
||||
if err != nil {
|
||||
appLogger.Error("Failed to list ollama models", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
appLogger.Info("Ollama models", map[string]interface{}{"models": aiModels})
|
||||
|
||||
// Initialize Worker
|
||||
bootstrap.InitWorker(*config, mongoManager, appLogger, notificationService, ollamaSDK)
|
||||
|
||||
// Set up signal handling for graceful shutdown
|
||||
signalChan := make(chan os.Signal, 1)
|
||||
signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
|
||||
|
||||
// Wait for shutdown signal
|
||||
<-signalChan
|
||||
appLogger.Info("Received shutdown signal, stopping scheduler...", map[string]interface{}{})
|
||||
|
||||
appLogger.Info("Worker application stopped", map[string]interface{}{})
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package workers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"tm/internal/notice"
|
||||
"tm/internal/tender"
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
"tm/pkg/notification"
|
||||
"tm/pkg/ollama"
|
||||
)
|
||||
|
||||
type NoticeWorker struct {
|
||||
Mongo *mongo.ConnectionManager
|
||||
Logger logger.Logger
|
||||
Notify notification.SDK
|
||||
Ollama *ollama.SDK
|
||||
NoticeRepo notice.Repository
|
||||
TenderRepo tender.TenderRepository
|
||||
}
|
||||
|
||||
func NewNoticeWorker(mongo *mongo.ConnectionManager, logger logger.Logger, notify notification.SDK, noticeRepo notice.Repository, tenderRepo tender.TenderRepository, ollama *ollama.SDK) *NoticeWorker {
|
||||
return &NoticeWorker{
|
||||
Mongo: mongo,
|
||||
Logger: logger,
|
||||
Notify: notify,
|
||||
Ollama: ollama,
|
||||
NoticeRepo: noticeRepo,
|
||||
TenderRepo: tenderRepo,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *NoticeWorker) Run() {
|
||||
w.Logger.Info("Notice worker started", map[string]interface{}{})
|
||||
|
||||
limit := 10
|
||||
skip := 0
|
||||
for {
|
||||
notices, _, err := w.NoticeRepo.GetUnProcessedNotices(context.Background(), limit, skip)
|
||||
if err != nil {
|
||||
w.Logger.Error("Failed to get notices", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
if len(notices) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
for _, n := range notices {
|
||||
w.Logger.Info("Notice", map[string]interface{}{
|
||||
"notice": n.ID.Hex(),
|
||||
})
|
||||
|
||||
t, err := w.TenderRepo.GetByNoticePublicationID(context.Background(), n.NoticePublicationID)
|
||||
if err == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
err = w.ToTender(&n, t)
|
||||
if err != nil {
|
||||
w.Logger.Error("Failed to create tender", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
if t.ID.IsZero() {
|
||||
err = w.TenderRepo.Create(context.Background(), t)
|
||||
if err != nil {
|
||||
w.Logger.Error("Failed to create tender", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
} else {
|
||||
err = w.TenderRepo.Update(context.Background(), t)
|
||||
if err != nil {
|
||||
w.Logger.Error("Failed to update tender", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
skip += limit
|
||||
}
|
||||
}
|
||||
|
||||
func (w *NoticeWorker) ToTender(n *notice.Notice, t *tender.Tender) error {
|
||||
// Buyer Organization
|
||||
buyerOrg := &tender.Organization{}
|
||||
if n.BuyerOrganization != nil {
|
||||
buyerOrg = &tender.Organization{
|
||||
Name: n.BuyerOrganization.Name,
|
||||
CompanyID: n.BuyerOrganization.CompanyID,
|
||||
WebsiteURI: n.BuyerOrganization.WebsiteURI,
|
||||
ContactName: n.BuyerOrganization.ContactName,
|
||||
ContactTelephone: n.BuyerOrganization.ContactTelephone,
|
||||
ContactEmail: n.BuyerOrganization.ContactEmail,
|
||||
ContactFax: n.BuyerOrganization.ContactFax,
|
||||
Role: n.BuyerOrganization.Role,
|
||||
Address: tender.Address{
|
||||
StreetName: n.BuyerOrganization.Address.StreetName,
|
||||
CityName: n.BuyerOrganization.Address.CityName,
|
||||
PostalZone: n.BuyerOrganization.Address.PostalZone,
|
||||
CountrySubentityCode: n.BuyerOrganization.Address.CountrySubentityCode,
|
||||
Department: n.BuyerOrganization.Address.Department,
|
||||
Region: n.BuyerOrganization.Address.Region,
|
||||
CountryCode: n.BuyerOrganization.Address.CountryCode,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Review Organization
|
||||
reviewOrg := &tender.Organization{}
|
||||
if n.ReviewOrganization != nil {
|
||||
reviewOrg = &tender.Organization{
|
||||
Name: n.ReviewOrganization.Name,
|
||||
CompanyID: n.ReviewOrganization.CompanyID,
|
||||
WebsiteURI: n.ReviewOrganization.WebsiteURI,
|
||||
ContactName: n.ReviewOrganization.ContactName,
|
||||
ContactTelephone: n.ReviewOrganization.ContactTelephone,
|
||||
ContactEmail: n.ReviewOrganization.ContactEmail,
|
||||
ContactFax: n.ReviewOrganization.ContactFax,
|
||||
Role: n.ReviewOrganization.Role,
|
||||
Address: tender.Address{
|
||||
StreetName: n.ReviewOrganization.Address.StreetName,
|
||||
CityName: n.ReviewOrganization.Address.CityName,
|
||||
PostalZone: n.ReviewOrganization.Address.PostalZone,
|
||||
CountrySubentityCode: n.ReviewOrganization.Address.CountrySubentityCode,
|
||||
Department: n.ReviewOrganization.Address.Department,
|
||||
Region: n.ReviewOrganization.Address.Region,
|
||||
CountryCode: n.ReviewOrganization.Address.CountryCode,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Organizations
|
||||
organizations := make([]tender.Organization, len(n.Organizations))
|
||||
for i, org := range n.Organizations {
|
||||
organizations[i] = tender.Organization{
|
||||
Name: org.Name,
|
||||
CompanyID: org.CompanyID,
|
||||
WebsiteURI: org.WebsiteURI,
|
||||
ContactName: org.ContactName,
|
||||
ContactTelephone: org.ContactTelephone,
|
||||
ContactEmail: org.ContactEmail,
|
||||
ContactFax: org.ContactFax,
|
||||
Role: org.Role,
|
||||
Address: tender.Address{
|
||||
StreetName: org.Address.StreetName,
|
||||
CityName: org.Address.CityName,
|
||||
PostalZone: org.Address.PostalZone,
|
||||
CountrySubentityCode: org.Address.CountrySubentityCode,
|
||||
Department: org.Address.Department,
|
||||
Region: org.Address.Region,
|
||||
CountryCode: org.Address.CountryCode,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Selection Criteria
|
||||
selectionCriteria := make([]tender.SelectionCriterion, len(n.SelectionCriteria))
|
||||
for i, criterion := range n.SelectionCriteria {
|
||||
selectionCriteria[i] = tender.SelectionCriterion{
|
||||
TypeCode: criterion.TypeCode,
|
||||
Description: criterion.Description,
|
||||
LanguageID: criterion.LanguageID,
|
||||
}
|
||||
}
|
||||
|
||||
// Winning Tenderer
|
||||
winningTenderer := &tender.Organization{}
|
||||
if n.WinningTenderer != nil {
|
||||
winningTenderer = &tender.Organization{
|
||||
Name: n.WinningTenderer.Name,
|
||||
CompanyID: n.WinningTenderer.CompanyID,
|
||||
WebsiteURI: n.WinningTenderer.WebsiteURI,
|
||||
ContactName: n.WinningTenderer.ContactName,
|
||||
ContactTelephone: n.WinningTenderer.ContactTelephone,
|
||||
ContactEmail: n.WinningTenderer.ContactEmail,
|
||||
ContactFax: n.WinningTenderer.ContactFax,
|
||||
Role: n.WinningTenderer.Role,
|
||||
Address: tender.Address{
|
||||
StreetName: n.WinningTenderer.Address.StreetName,
|
||||
CityName: n.WinningTenderer.Address.CityName,
|
||||
PostalZone: n.WinningTenderer.Address.PostalZone,
|
||||
CountrySubentityCode: n.WinningTenderer.Address.CountrySubentityCode,
|
||||
Department: n.WinningTenderer.Address.Department,
|
||||
Region: n.WinningTenderer.Address.Region,
|
||||
CountryCode: n.WinningTenderer.Address.CountryCode,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Modifications
|
||||
modifications := make([]tender.TenderModification, len(n.Modifications))
|
||||
for i, modification := range n.Modifications {
|
||||
modifications[i] = tender.TenderModification{
|
||||
ModificationDate: modification.ModificationDate,
|
||||
ModificationReason: modification.ModificationReason,
|
||||
Description: modification.Description,
|
||||
LanguageID: modification.LanguageID,
|
||||
}
|
||||
}
|
||||
|
||||
// Awarded Entities
|
||||
awardedEntities := make([]tender.Awarded, len(n.AwardedEntities))
|
||||
for i, awarded := range n.AwardedEntities {
|
||||
awardedEntities[i] = tender.Awarded{
|
||||
Name: awarded.Name,
|
||||
Address: awarded.Address,
|
||||
Country: awarded.Country,
|
||||
Amount: awarded.Amount,
|
||||
Currency: awarded.Currency,
|
||||
Share: awarded.Share,
|
||||
AwardDate: awarded.AwardDate,
|
||||
ContractID: awarded.ContractID,
|
||||
TenderID: awarded.TenderID,
|
||||
LotID: awarded.LotID,
|
||||
CompanyID: awarded.CompanyID,
|
||||
OrganizationID: awarded.OrganizationID,
|
||||
}
|
||||
}
|
||||
|
||||
t.NoticePublicationID = n.NoticePublicationID
|
||||
t.Title = n.Title
|
||||
t.Description = n.Description
|
||||
t.ProcurementTypeCode = n.ProcurementTypeCode
|
||||
t.ProcedureCode = n.ProcedureCode
|
||||
t.EstimatedValue = n.EstimatedValue
|
||||
t.Currency = n.Currency
|
||||
t.TenderDeadline = n.TenderDeadline
|
||||
t.SubmissionURL = n.SubmissionURL
|
||||
t.BuyerOrganization = buyerOrg
|
||||
t.ReviewOrganization = reviewOrg
|
||||
t.Organizations = organizations
|
||||
t.SelectionCriteria = selectionCriteria
|
||||
t.WinningTenderer = winningTenderer
|
||||
t.Modifications = modifications
|
||||
t.AwardedEntities = awardedEntities
|
||||
t.NoticeTypeCode = n.NoticeTypeCode
|
||||
t.NoticeSubTypeCode = n.NoticeSubTypeCode
|
||||
t.NoticeLanguageCode = n.NoticeLanguageCode
|
||||
t.IssueDate = n.IssueDate
|
||||
t.IssueTime = n.IssueTime
|
||||
t.PublicationDate = n.PublicationDate
|
||||
t.SubmissionDeadline = n.SubmissionDeadline
|
||||
t.ApplicationDeadline = n.ApplicationDeadline
|
||||
t.GazetteID = n.GazetteID
|
||||
t.PlaceOfPerformance = n.PlaceOfPerformance
|
||||
t.CountryCode = n.CountryCode
|
||||
t.RegionCode = n.RegionCode
|
||||
t.CityName = n.CityName
|
||||
t.PostalCode = n.PostalCode
|
||||
t.DocumentURI = n.DocumentURI
|
||||
t.TenderURL = n.TenderURL
|
||||
t.MainClassification = n.MainClassification
|
||||
t.AdditionalClassifications = n.AdditionalClassifications
|
||||
t.OfficialLanguages = n.OfficialLanguages
|
||||
t.CancellationReason = n.CancellationReason
|
||||
t.CancellationDate = n.CancellationDate
|
||||
t.AwardDate = n.AwardDate
|
||||
t.AwardedValue = n.AwardedValue
|
||||
t.ContractNumber = n.ContractNumber
|
||||
t.SuspensionReason = n.SuspensionReason
|
||||
t.SuspensionDate = n.SuspensionDate
|
||||
t.WinningTenderer = winningTenderer
|
||||
t.Modifications = modifications
|
||||
t.AwardedEntities = awardedEntities
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
package notice
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"tm/pkg/mongo"
|
||||
)
|
||||
|
||||
// Tender represents a tender/contract notice entity
|
||||
type Notice struct {
|
||||
mongo.Model `bson:",inline"`
|
||||
ContractNoticeID string `bson:"contract_notice_id" json:"contract_notice_id"`
|
||||
NoticePublicationID string `bson:"notice_publication_id" json:"notice_publication_id"`
|
||||
ContractFolderID string `bson:"contract_folder_id" json:"contract_folder_id"`
|
||||
NoticeTypeCode string `bson:"notice_type_code" json:"notice_type_code"`
|
||||
NoticeSubTypeCode string `bson:"notice_sub_type_code" json:"notice_sub_type_code"`
|
||||
NoticeLanguageCode string `bson:"notice_language_code" json:"notice_language_code"`
|
||||
IssueDate int64 `bson:"issue_date" json:"issue_date"`
|
||||
IssueTime int64 `bson:"issue_time" json:"issue_time"`
|
||||
TenderDeadline int64 `bson:"tender_deadline" json:"tender_deadline"`
|
||||
PublicationDate int64 `bson:"publication_date" json:"publication_date"`
|
||||
SubmissionDeadline int64 `bson:"submission_deadline" json:"submission_deadline"`
|
||||
ApplicationDeadline int64 `bson:"application_deadline" json:"application_deadline"`
|
||||
GazetteID string `bson:"gazette_id" json:"gazette_id"`
|
||||
Title string `bson:"title" json:"title"`
|
||||
Description string `bson:"description" json:"description"`
|
||||
ProcurementTypeCode string `bson:"procurement_type_code" json:"procurement_type_code"`
|
||||
ProcedureCode string `bson:"procedure_code" json:"procedure_code"`
|
||||
MainClassification string `bson:"main_classification" json:"main_classification"`
|
||||
AdditionalClassifications []string `bson:"additional_classifications" json:"additional_classifications"`
|
||||
EstimatedValue float64 `bson:"estimated_value" json:"estimated_value"`
|
||||
Currency string `bson:"currency" json:"currency"`
|
||||
Duration string `bson:"duration" json:"duration"`
|
||||
DurationUnit string `bson:"duration_unit" json:"duration_unit"`
|
||||
PlaceOfPerformance string `bson:"place_of_performance" json:"place_of_performance"`
|
||||
CountryCode string `bson:"country_code" json:"country_code"`
|
||||
RegionCode string `bson:"region_code" json:"region_code"`
|
||||
CityName string `bson:"city_name" json:"city_name"`
|
||||
PostalCode string `bson:"postal_code" json:"postal_code"`
|
||||
DocumentURI string `bson:"document_uri" json:"document_uri"`
|
||||
TenderURL string `bson:"tender_url" json:"tender_url"`
|
||||
SubmissionURL string `bson:"submission_url" json:"submission_url"`
|
||||
BuyerOrganization *Organization `bson:"buyer_organization" json:"buyer_organization"`
|
||||
ReviewOrganization *Organization `bson:"review_organization,omitempty" json:"review_organization,omitempty"`
|
||||
Organizations []Organization `bson:"organizations" json:"organizations"`
|
||||
SelectionCriteria []SelectionCriterion `bson:"selection_criteria" json:"selection_criteria"`
|
||||
OfficialLanguages []string `bson:"official_languages" json:"official_languages"`
|
||||
Status TenderStatus `bson:"status" json:"status"`
|
||||
Source TenderSource `bson:"source" json:"source"`
|
||||
SourceFileURL string `bson:"source_file_url" json:"source_file_url"`
|
||||
SourceFileName string `bson:"source_file_name" json:"source_file_name"`
|
||||
ProcessingMetadata ProcessingMetadata `bson:"processing_metadata" json:"processing_metadata"`
|
||||
TenderID string `bson:"tender_id" json:"tender_id"` // Implement logic to generate a unique tender_id using the following format: {Source}{BUYER_CODE}{TED_CODE}{DATE}{LOCATION_CODE}
|
||||
|
||||
// Status-specific fields
|
||||
CancellationReason string `bson:"cancellation_reason,omitempty" json:"cancellation_reason,omitempty"`
|
||||
CancellationDate int64 `bson:"cancellation_date,omitempty" json:"cancellation_date,omitempty"`
|
||||
AwardDate int64 `bson:"award_date,omitempty" json:"award_date,omitempty"`
|
||||
AwardedValue float64 `bson:"awarded_value,omitempty" json:"awarded_value,omitempty"`
|
||||
WinningTenderer *Organization `bson:"winning_tenderer,omitempty" json:"winning_tenderer,omitempty"`
|
||||
ContractNumber string `bson:"contract_number,omitempty" json:"contract_number,omitempty"`
|
||||
SuspensionReason string `bson:"suspension_reason,omitempty" json:"suspension_reason,omitempty"`
|
||||
SuspensionDate int64 `bson:"suspension_date,omitempty" json:"suspension_date,omitempty"`
|
||||
Modifications []TenderModification `bson:"modifications,omitempty" json:"modifications,omitempty"`
|
||||
AwardedEntities []Awarded `bson:"awarded,omitempty" json:"awarded,omitempty"`
|
||||
}
|
||||
|
||||
// Organization represents organization information
|
||||
type Organization struct {
|
||||
ID string `bson:"id" json:"id"`
|
||||
Name string `bson:"name" json:"name"`
|
||||
CompanyID string `bson:"company_id,omitempty" json:"company_id,omitempty"`
|
||||
WebsiteURI string `bson:"website_uri,omitempty" json:"website_uri,omitempty"`
|
||||
ContactName string `bson:"contact_name,omitempty" json:"contact_name,omitempty"`
|
||||
ContactTelephone string `bson:"contact_telephone,omitempty" json:"contact_telephone,omitempty"`
|
||||
ContactEmail string `bson:"contact_email,omitempty" json:"contact_email,omitempty"`
|
||||
ContactFax string `bson:"contact_fax,omitempty" json:"contact_fax,omitempty"`
|
||||
Address Address `bson:"address" json:"address"`
|
||||
Role string `bson:"role,omitempty" json:"role,omitempty"`
|
||||
}
|
||||
|
||||
// Address represents address information
|
||||
type Address struct {
|
||||
StreetName string `bson:"street_name,omitempty" json:"street_name,omitempty"`
|
||||
CityName string `bson:"city_name,omitempty" json:"city_name,omitempty"`
|
||||
PostalZone string `bson:"postal_zone,omitempty" json:"postal_zone,omitempty"`
|
||||
CountrySubentityCode string `bson:"country_subentity_code,omitempty" json:"country_subentity_code,omitempty"`
|
||||
Department string `bson:"department,omitempty" json:"department,omitempty"`
|
||||
Region string `bson:"region,omitempty" json:"region,omitempty"`
|
||||
CountryCode string `bson:"country_code" json:"country_code"`
|
||||
}
|
||||
|
||||
// SelectionCriterion represents selection criteria
|
||||
type SelectionCriterion struct {
|
||||
TypeCode string `bson:"type_code" json:"type_code"`
|
||||
Description string `bson:"description,omitempty" json:"description,omitempty"`
|
||||
LanguageID string `bson:"language_id,omitempty" json:"language_id,omitempty"`
|
||||
}
|
||||
|
||||
// ProcessingMetadata contains metadata about how the tender was processed
|
||||
type ProcessingMetadata struct {
|
||||
ScrapedAt int64 `bson:"scraped_at" json:"scraped_at"`
|
||||
ProcessedAt int64 `bson:"processed_at" json:"processed_at"`
|
||||
ProcessingVersion string `bson:"processing_version" json:"processing_version"`
|
||||
ParsingErrors []string `bson:"parsing_errors,omitempty" json:"parsing_errors,omitempty"`
|
||||
ValidationErrors []string `bson:"validation_errors,omitempty" json:"validation_errors,omitempty"`
|
||||
EnrichmentData map[string]string `bson:"enrichment_data,omitempty" json:"enrichment_data,omitempty"`
|
||||
TranslatedData map[string]string `bson:"translated_data,omitempty" json:"translated_data,omitempty"`
|
||||
TranslatedAt int64 `bson:"translated_at,omitempty" json:"translated_at,omitempty"`
|
||||
Processed bool `bson:"processed" json:"processed"`
|
||||
}
|
||||
|
||||
// TenderModification represents a modification made to a tender
|
||||
type TenderModification struct {
|
||||
ModificationDate int64 `bson:"modification_date" json:"modification_date"` // Unix milliseconds
|
||||
ModificationReason string `bson:"modification_reason" json:"modification_reason"`
|
||||
Description string `bson:"description,omitempty" json:"description,omitempty"`
|
||||
LanguageID string `bson:"language_id,omitempty" json:"language_id,omitempty"`
|
||||
}
|
||||
|
||||
// AwardedEntity represents an entity that has been awarded a contract
|
||||
type Awarded struct {
|
||||
Name string `bson:"name" json:"name"`
|
||||
Address string `bson:"address,omitempty" json:"address,omitempty"`
|
||||
Country string `bson:"country,omitempty" json:"country,omitempty"`
|
||||
Amount float64 `bson:"amount" json:"amount"`
|
||||
Currency string `bson:"currency,omitempty" json:"currency,omitempty"`
|
||||
Share float64 `bson:"share,omitempty" json:"share,omitempty"` // Percentage share for multiple winners
|
||||
AwardDate int64 `bson:"award_date" json:"award_date"` // Unix milliseconds
|
||||
ContractID string `bson:"contract_id,omitempty" json:"contract_id,omitempty"` // Contract reference ID
|
||||
TenderID string `bson:"tender_id,omitempty" json:"tender_id,omitempty"` // Tender reference ID
|
||||
LotID string `bson:"lot_id,omitempty" json:"lot_id,omitempty"` // Lot reference ID
|
||||
CompanyID string `bson:"company_id,omitempty" json:"company_id,omitempty"` // Company registration ID
|
||||
OrganizationID string `bson:"organization_id,omitempty" json:"organization_id,omitempty"` // Organization reference ID
|
||||
}
|
||||
|
||||
// TenderStatus represents the status of a tender
|
||||
type TenderStatus string
|
||||
|
||||
const (
|
||||
TenderStatusActive TenderStatus = "active"
|
||||
TenderStatusExpired TenderStatus = "expired"
|
||||
TenderStatusCancelled TenderStatus = "cancelled"
|
||||
TenderStatusAwarded TenderStatus = "awarded"
|
||||
TenderStatusDraft TenderStatus = "draft"
|
||||
TenderStatusClosed TenderStatus = "closed"
|
||||
TenderStatusModified TenderStatus = "modified"
|
||||
TenderStatusSuspended TenderStatus = "suspended"
|
||||
TenderStatusPublished TenderStatus = "published"
|
||||
)
|
||||
|
||||
// TenderSource represents the source of tender data
|
||||
type TenderSource string
|
||||
|
||||
const (
|
||||
TenderSourceTEDScraper TenderSource = "ted_scraper"
|
||||
TenderSourceManual TenderSource = "manual"
|
||||
TenderSourceAPI TenderSource = "api"
|
||||
TenderSourceBulkImport TenderSource = "bulk_import"
|
||||
)
|
||||
|
||||
// TenderStatistics represents tender statistics
|
||||
type TenderStatistics struct {
|
||||
TotalTenders int64 `json:"total_tenders"`
|
||||
ActiveTenders int64 `json:"active_tenders"`
|
||||
ExpiredTenders int64 `json:"expired_tenders"`
|
||||
TendersByCountry map[string]int64 `json:"tenders_by_country"`
|
||||
TendersByType map[string]int64 `json:"tenders_by_type"`
|
||||
TendersByClassification map[string]int64 `json:"tenders_by_classification"`
|
||||
AverageEstimatedValue float64 `json:"average_estimated_value"`
|
||||
TotalEstimatedValue float64 `json:"total_estimated_value"`
|
||||
LastUpdated int64 `json:"last_updated"` // Unix milliseconds
|
||||
}
|
||||
|
||||
// GetID returns the tender ID
|
||||
func (t *Notice) GetID() string {
|
||||
return t.ID.Hex()
|
||||
}
|
||||
|
||||
// IsActive returns true if the tender is active
|
||||
func (t *Notice) IsActive() bool {
|
||||
return t.Status == TenderStatusActive
|
||||
}
|
||||
|
||||
// IsExpired returns true if the tender deadline has passed
|
||||
func (t *Notice) IsExpired() bool {
|
||||
if t.Status == TenderStatusExpired {
|
||||
return true
|
||||
}
|
||||
|
||||
if t.TenderDeadline == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if deadline (in milliseconds) is in the past
|
||||
now := time.Now().Unix()
|
||||
return t.TenderDeadline < now
|
||||
}
|
||||
|
||||
// GetTenderURL returns the TED tender URL
|
||||
func (t *Notice) GetTenderURL() string {
|
||||
if t.TenderURL != "" {
|
||||
return t.TenderURL
|
||||
}
|
||||
|
||||
if t.NoticePublicationID != "" {
|
||||
return generateTenderURL(t.NoticePublicationID)
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// generateTenderURL generates the TED tender URL from NoticePublicationID
|
||||
func generateTenderURL(noticePublicationID string) string {
|
||||
// Remove leading zeros and format for URL
|
||||
// Example: 00522942-2025 -> https://ted.europa.eu/en/notice/-/detail/522942-2025
|
||||
parts := strings.Split(noticePublicationID, "-")
|
||||
if len(parts) != 2 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Convert to int and back to string to remove leading zeros
|
||||
if num, err := strconv.Atoi(parts[0]); err == nil {
|
||||
cleanID := fmt.Sprintf("%d-%s", num, parts[1])
|
||||
return fmt.Sprintf("https://ted.europa.eu/en/notice/-/detail/%s", cleanID)
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetMainOrganization returns the buyer organization
|
||||
func (t *Notice) GetMainOrganization() *Organization {
|
||||
return t.BuyerOrganization
|
||||
}
|
||||
|
||||
// HasEstimatedValue returns true if the tender has an estimated value
|
||||
func (t *Notice) HasEstimatedValue() bool {
|
||||
return t.EstimatedValue > 0
|
||||
}
|
||||
|
||||
// GetFormattedEstimatedValue returns formatted estimated value with currency
|
||||
func (t *Notice) GetFormattedEstimatedValue() string {
|
||||
if !t.HasEstimatedValue() {
|
||||
return ""
|
||||
}
|
||||
|
||||
if t.Currency != "" {
|
||||
return fmt.Sprintf("%.2f %s", t.EstimatedValue, t.Currency)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%.2f", t.EstimatedValue)
|
||||
}
|
||||
|
||||
// UpdateStatus updates the tender status and updated timestamp
|
||||
func (t *Notice) UpdateStatus(status TenderStatus) {
|
||||
t.Status = status
|
||||
t.UpdatedAt = time.Now().Unix()
|
||||
}
|
||||
|
||||
// GetDaysUntilDeadline returns the number of days until the tender deadline
|
||||
func (t *Notice) GetDaysUntilDeadline() int {
|
||||
if t.TenderDeadline == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
now := time.Now().Unix()
|
||||
if t.TenderDeadline <= now {
|
||||
return 0 // Already expired
|
||||
}
|
||||
|
||||
diffMillis := t.TenderDeadline - now
|
||||
diffDays := diffMillis / (24 * 60 * 60 * 1000)
|
||||
return int(diffDays)
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package notice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
"tm/pkg/logger"
|
||||
orm "tm/pkg/mongo"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
)
|
||||
|
||||
// NoticeRepository interface defines notice data access methods
|
||||
type Repository interface {
|
||||
Import(ctx context.Context, notice *Notice) error
|
||||
BulkImport(ctx context.Context, notices []Notice) error
|
||||
GetByID(ctx context.Context, id string) (*Notice, error)
|
||||
GetByContractNoticeID(ctx context.Context, contractNoticeID string) (*Notice, error)
|
||||
GetUnProcessedNotices(ctx context.Context, limit int, skip int) ([]Notice, int64, error)
|
||||
Delete(ctx context.Context, id string) error
|
||||
}
|
||||
|
||||
func collectionName() string {
|
||||
return "notices"
|
||||
}
|
||||
|
||||
// noticeRepository implements NoticeRepository interface using MongoDB ORM
|
||||
type noticeRepository struct {
|
||||
ormRepo orm.Repository[Notice]
|
||||
mongoManager *orm.ConnectionManager
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
// NewRepository creates a new notice repository
|
||||
func NewRepository(mongoManager *orm.ConnectionManager, logger logger.Logger) Repository {
|
||||
// Create indexes for notices collection
|
||||
noticeIndexes := []orm.Index{}
|
||||
|
||||
// Create indexes
|
||||
if err := mongoManager.CreateIndexes(collectionName(), noticeIndexes); err != nil {
|
||||
logger.Warn("Failed to create notice indexes", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
// Create ORM repositories
|
||||
noticeOrmRepo := orm.NewRepository[Notice](mongoManager.GetCollection(collectionName()), logger)
|
||||
|
||||
return ¬iceRepository{
|
||||
ormRepo: noticeOrmRepo,
|
||||
mongoManager: mongoManager,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// Import imports a new notice
|
||||
func (r *noticeRepository) Import(ctx context.Context, notice *Notice) error {
|
||||
now := time.Now().Unix()
|
||||
notice.SetCreatedAt(now)
|
||||
|
||||
err := r.ormRepo.Create(ctx, notice)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to create notice", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
"notice_publication_id": notice.NoticePublicationID,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
r.logger.Info("Notice created successfully", map[string]interface{}{
|
||||
"notice_id": notice.ID,
|
||||
"notice_publication_id": notice.NoticePublicationID,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// BulkImport imports a new notice
|
||||
func (r *noticeRepository) BulkImport(ctx context.Context, notices []Notice) error {
|
||||
err := r.ormRepo.CreateMany(ctx, notices)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to bulk import notices", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetByID retrieves a notice by ID
|
||||
func (r *noticeRepository) GetByID(ctx context.Context, id string) (*Notice, error) {
|
||||
notice, err := r.ormRepo.FindByID(ctx, id)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to get notice by ID", map[string]interface{}{
|
||||
"notice_id": id,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return notice, nil
|
||||
}
|
||||
|
||||
// GetByContractNoticeID retrieves a notice by contract notice ID
|
||||
func (r *noticeRepository) GetByContractNoticeID(ctx context.Context, contractNoticeID string) (*Notice, error) {
|
||||
filter := bson.M{"contract_notice_id": contractNoticeID}
|
||||
pagination := orm.Pagination{Limit: 1}
|
||||
|
||||
result, err := r.ormRepo.FindAll(ctx, filter, pagination)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to get notice by contract notice ID", map[string]interface{}{
|
||||
"contract_notice_id": contractNoticeID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(result.Items) == 0 {
|
||||
return nil, orm.ErrDocumentNotFound
|
||||
}
|
||||
|
||||
return &result.Items[0], nil
|
||||
}
|
||||
|
||||
// GetUnProcessedNotices retrieves unprocessed notices
|
||||
|
||||
func (r *noticeRepository) GetUnProcessedNotices(ctx context.Context, limit int, skip int) ([]Notice, int64, error) {
|
||||
filter := bson.M{"processing_metadata.processed": false}
|
||||
|
||||
pagination := orm.Pagination{
|
||||
Limit: limit,
|
||||
Skip: skip,
|
||||
}
|
||||
|
||||
result, err := r.ormRepo.FindAll(ctx, filter, pagination)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to get unprocessed notices", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return result.Items, result.TotalCount, nil
|
||||
}
|
||||
|
||||
// Update updates an existing notice
|
||||
func (r *noticeRepository) Update(ctx context.Context, notice *Notice) error {
|
||||
notice.UpdatedAt = time.Now().Unix()
|
||||
|
||||
err := r.ormRepo.Update(ctx, notice)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to update notice", map[string]interface{}{
|
||||
"notice_id": notice.ID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
r.logger.Info("Notice updated successfully", map[string]interface{}{
|
||||
"notice_id": notice.ID,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete deletes a notice by ID
|
||||
func (r *noticeRepository) Delete(ctx context.Context, id string) error {
|
||||
err := r.ormRepo.Delete(ctx, id)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to delete notice", map[string]interface{}{
|
||||
"notice_id": id,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
r.logger.Info("Notice deleted successfully", map[string]interface{}{
|
||||
"notice_id": id,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package notice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"tm/pkg/logger"
|
||||
)
|
||||
|
||||
// Service defines business logic for notice operations
|
||||
type Service interface {
|
||||
Import(ctx context.Context, form *Notice) error
|
||||
BulkImport(ctx context.Context, form []Notice) error
|
||||
}
|
||||
|
||||
// noticeService implements the Service interface
|
||||
type noticeService struct {
|
||||
repository Repository
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
// NewService creates a new notice service
|
||||
func NewService(repository Repository, logger logger.Logger) Service {
|
||||
return ¬iceService{
|
||||
repository: repository,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// Import imports a new notice
|
||||
func (s *noticeService) Import(ctx context.Context, form *Notice) error {
|
||||
return s.repository.Import(ctx, form)
|
||||
}
|
||||
|
||||
// BulkImport imports a new notice
|
||||
func (s *noticeService) BulkImport(ctx context.Context, form []Notice) error {
|
||||
return s.repository.BulkImport(ctx, form)
|
||||
}
|
||||
@@ -74,6 +74,21 @@ type NotificationConfig struct {
|
||||
UserAgent string `env:"NOTIFICATION_USER_AGENT"`
|
||||
}
|
||||
|
||||
// OllamaConfig represents the configuration for the Ollama SDK
|
||||
type OllamaConfig struct {
|
||||
BaseURL string `env:"OLLAMA_BASE_URL" default:"http://localhost:11434"`
|
||||
Timeout time.Duration `env:"OLLAMA_TIMEOUT" default:"300s"`
|
||||
RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS" default:"3"`
|
||||
RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY" default:"2s"`
|
||||
EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING" default:"true"`
|
||||
UserAgent string `env:"OLLAMA_USER_AGENT" default:"TenderManagement-OllamaSDK/1.0"`
|
||||
DefaultModel string `env:"OLLAMA_DEFAULT_MODEL" default:"llama2"`
|
||||
DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE" default:"0.7"`
|
||||
DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P" default:"0.9"`
|
||||
DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS" default:"1000"`
|
||||
EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING" default:"false"`
|
||||
}
|
||||
|
||||
// LoadConfig loads configuration with priority: OS environment variables > .env file
|
||||
// OS environment variables take precedence over .env file values
|
||||
func LoadConfig[T any](path string, config T) (T, error) {
|
||||
|
||||
@@ -47,6 +47,9 @@ type Repository[T any] interface {
|
||||
// Create inserts a new document
|
||||
Create(ctx context.Context, model *T) error
|
||||
|
||||
// CreateMany inserts multiple documents
|
||||
CreateMany(ctx context.Context, models []T) error
|
||||
|
||||
// Update updates an existing document
|
||||
Update(ctx context.Context, model *T) error
|
||||
|
||||
@@ -239,6 +242,32 @@ func (r *repository[T]) Create(ctx context.Context, model *T) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateMany inserts multiple documents
|
||||
func (r *repository[T]) CreateMany(ctx context.Context, models []T) error {
|
||||
// Set timestamps if the model implements Timestampable
|
||||
for _, v := range models {
|
||||
if timestampable, ok := any(v).(Timestampable); ok {
|
||||
now := time.Now().Unix()
|
||||
timestampable.SetCreatedAt(now)
|
||||
timestampable.SetUpdatedAt(now)
|
||||
}
|
||||
}
|
||||
|
||||
result, err := r.collection.InsertMany(ctx, models)
|
||||
if err != nil {
|
||||
r.logger.Error("Failed to create multiple documents", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
return fmt.Errorf("failed to create multiple documents: %w", err)
|
||||
}
|
||||
|
||||
r.logger.Info("Documents created successfully", map[string]interface{}{
|
||||
"ids": result.InsertedIDs,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Update updates an existing document
|
||||
func (r *repository[T]) Update(ctx context.Context, model *T) error {
|
||||
// Set updated timestamp if the model implements Timestampable
|
||||
|
||||
@@ -0,0 +1,394 @@
|
||||
# Ollama SDK
|
||||
|
||||
A comprehensive Go SDK for interacting with Ollama AI models. This SDK provides a clean, type-safe interface for text generation, chat conversations, embeddings, and model management.
|
||||
|
||||
## Features
|
||||
|
||||
- **Text Generation**: Generate text using various AI models
|
||||
- **Chat Conversations**: Interactive chat with support for system prompts and conversation history
|
||||
- **Embeddings**: Generate embeddings for text similarity and search
|
||||
- **Model Management**: Pull, list, delete, and inspect models
|
||||
- **Streaming Support**: Real-time streaming of generated text and chat responses
|
||||
- **Automatic Retries**: Configurable retry mechanism with exponential backoff
|
||||
- **Structured Logging**: Built-in logging with structured fields
|
||||
- **Fluent API**: Builder pattern for constructing complex chat conversations
|
||||
- **Type Safety**: Full type safety with comprehensive validation
|
||||
- **Error Handling**: Detailed error types and proper error wrapping
|
||||
- **Configuration**: Flexible configuration via environment variables or code
|
||||
- **Health Checks**: Built-in health checking capabilities
|
||||
|
||||
## Installation
|
||||
|
||||
The SDK is already part of your `tm` module. Import it as:
|
||||
```go
|
||||
import "tm/pkg/ollama"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Basic Setup
|
||||
|
||||
```go
|
||||
// 1. Configuration
|
||||
config := ollama.DefaultConfig()
|
||||
config.BaseURL = "http://localhost:11434" // Your Ollama server URL
|
||||
config.DefaultModel = "llama2" // Default model to use
|
||||
|
||||
// 2. Logger (using project's logger)
|
||||
loggerConfig := &logger.Config{Level: "info", Format: "json", Output: "stdout"}
|
||||
projectLogger := logger.NewLogger(loggerConfig)
|
||||
|
||||
// 3. Create SDK
|
||||
sdk, err := ollama.New(config, projectLogger)
|
||||
if err != nil {
|
||||
log.Fatal("Failed to create Ollama SDK:", err)
|
||||
}
|
||||
```
|
||||
|
||||
### Environment Variables
|
||||
|
||||
Set environment variables for automatic configuration:
|
||||
|
||||
```bash
|
||||
export OLLAMA_BASE_URL="http://localhost:11434"
|
||||
export OLLAMA_TIMEOUT="300s"
|
||||
export OLLAMA_RETRY_ATTEMPTS="3"
|
||||
export OLLAMA_RETRY_DELAY="2s"
|
||||
export OLLAMA_ENABLE_LOGGING="true"
|
||||
export OLLAMA_DEFAULT_MODEL="llama2"
|
||||
export OLLAMA_DEFAULT_TEMPERATURE="0.7"
|
||||
export OLLAMA_DEFAULT_TOP_P="0.9"
|
||||
export OLLAMA_DEFAULT_MAX_TOKENS="1000"
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Configuration
|
||||
|
||||
```go
|
||||
type Config struct {
|
||||
BaseURL string `env:"OLLAMA_BASE_URL"`
|
||||
Timeout time.Duration `env:"OLLAMA_TIMEOUT"`
|
||||
RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS"`
|
||||
RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY"`
|
||||
EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING"`
|
||||
UserAgent string `env:"OLLAMA_USER_AGENT"`
|
||||
DefaultModel string `env:"OLLAMA_DEFAULT_MODEL"`
|
||||
DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE"`
|
||||
DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P"`
|
||||
DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS"`
|
||||
EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING"`
|
||||
}
|
||||
```
|
||||
|
||||
### Creating SDK Instance
|
||||
|
||||
```go
|
||||
// With default configuration
|
||||
sdk, err := ollama.New(nil, logger)
|
||||
|
||||
// With custom configuration
|
||||
config := &ollama.Config{
|
||||
BaseURL: "http://localhost:11434",
|
||||
Timeout: 300 * time.Second,
|
||||
RetryAttempts: 3,
|
||||
RetryDelay: 2 * time.Second,
|
||||
EnableLogging: true,
|
||||
DefaultModel: "llama2",
|
||||
DefaultTemperature: 0.7,
|
||||
DefaultTopP: 0.9,
|
||||
DefaultMaxTokens: 1000,
|
||||
}
|
||||
sdk, err := ollama.New(config, logger)
|
||||
|
||||
// With custom HTTP client
|
||||
httpClient := &http.Client{Timeout: 10 * time.Second}
|
||||
sdk, err := ollama.NewWithClient(config, httpClient, logger)
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Simple Text Generation
|
||||
|
||||
```go
|
||||
ctx := context.Background()
|
||||
|
||||
// Quick generation with default model
|
||||
response, err := sdk.QuickGenerate(ctx, "Write a short story about a robot")
|
||||
if err != nil {
|
||||
log.Printf("Generation failed: %v", err)
|
||||
}
|
||||
fmt.Println(response)
|
||||
|
||||
// Generation with specific model
|
||||
response, err := sdk.QuickGenerateWithModel(ctx, "codellama", "Write a Python function to sort a list")
|
||||
```
|
||||
|
||||
### Advanced Text Generation
|
||||
|
||||
```go
|
||||
req := &ollama.GenerateRequest{
|
||||
Model: "llama2",
|
||||
Prompt: "Explain quantum computing in simple terms",
|
||||
Options: map[string]interface{}{
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"max_tokens": 500,
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := sdk.Generate(ctx, req)
|
||||
if err != nil {
|
||||
log.Printf("Generation failed: %v", err)
|
||||
}
|
||||
fmt.Printf("Generated: %s\n", resp.Response)
|
||||
```
|
||||
|
||||
### Chat Conversations
|
||||
|
||||
```go
|
||||
// Simple chat
|
||||
response, err := sdk.QuickChat(ctx, "What is the capital of France?")
|
||||
fmt.Println(response)
|
||||
|
||||
// Chat with system prompt
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a helpful coding assistant.",
|
||||
"How do I implement a binary search in Go?")
|
||||
```
|
||||
|
||||
### Advanced Chat with Conversation History
|
||||
|
||||
```go
|
||||
req := &ollama.ChatRequest{
|
||||
Model: "llama2",
|
||||
Messages: []ollama.Message{
|
||||
{
|
||||
Role: ollama.MessageRoleSystem,
|
||||
Content: "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
Role: ollama.MessageRoleUser,
|
||||
Content: "Hello, how are you?",
|
||||
},
|
||||
{
|
||||
Role: ollama.MessageRoleAssistant,
|
||||
Content: "I'm doing well, thank you! How can I help you today?",
|
||||
},
|
||||
{
|
||||
Role: ollama.MessageRoleUser,
|
||||
Content: "Can you help me with a coding problem?",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := sdk.Chat(ctx, req)
|
||||
fmt.Printf("Assistant: %s\n", resp.Message.Content)
|
||||
```
|
||||
|
||||
### Fluent Chat Builder
|
||||
|
||||
```go
|
||||
response, err := sdk.NewChatBuilder().
|
||||
SetModel("llama2").
|
||||
SetSystem("You are a helpful coding assistant.").
|
||||
AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?").
|
||||
SetTemperature(0.7).
|
||||
SetMaxTokens(500).
|
||||
Send(ctx)
|
||||
|
||||
fmt.Printf("Assistant: %s\n", response.Message.Content)
|
||||
|
||||
// Continue the conversation
|
||||
response, err = sdk.NewChatBuilder().
|
||||
SetModel("llama2").
|
||||
SetSystem("You are a helpful coding assistant.").
|
||||
AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?").
|
||||
SetTemperature(0.7).
|
||||
SetMaxTokens(500).
|
||||
Continue(ctx, "Can you show me an example with methods?")
|
||||
|
||||
fmt.Printf("Assistant: %s\n", response.Message.Content)
|
||||
```
|
||||
|
||||
### Streaming Responses
|
||||
|
||||
```go
|
||||
// Streaming text generation
|
||||
err := sdk.StreamGenerate(ctx, &ollama.GenerateRequest{
|
||||
Model: "llama2",
|
||||
Prompt: "Write a poem about programming",
|
||||
}, func(chunk string) error {
|
||||
fmt.Print(chunk)
|
||||
return nil
|
||||
})
|
||||
|
||||
// Streaming chat
|
||||
err := sdk.StreamChat(ctx, &ollama.ChatRequest{
|
||||
Model: "llama2",
|
||||
Messages: []ollama.Message{
|
||||
{
|
||||
Role: ollama.MessageRoleUser,
|
||||
Content: "Tell me a story",
|
||||
},
|
||||
},
|
||||
}, func(chunk string) error {
|
||||
fmt.Print(chunk)
|
||||
return nil
|
||||
})
|
||||
```
|
||||
|
||||
### Embeddings
|
||||
|
||||
```go
|
||||
// Generate embeddings
|
||||
embeddings, err := sdk.QuickEmbed(ctx, "This is some text to embed")
|
||||
if err != nil {
|
||||
log.Printf("Embedding failed: %v", err)
|
||||
}
|
||||
|
||||
// Advanced embedding request
|
||||
req := &ollama.EmbedRequest{
|
||||
Model: "llama2",
|
||||
Prompt: "Text to generate embeddings for",
|
||||
}
|
||||
|
||||
resp, err := sdk.Embed(ctx, req)
|
||||
fmt.Printf("Embedding dimension: %d\n", len(resp.Embedding))
|
||||
```
|
||||
|
||||
### Model Management
|
||||
|
||||
```go
|
||||
// List available models
|
||||
models, err := sdk.ListModels(ctx)
|
||||
for _, model := range models.Models {
|
||||
fmt.Printf("Model: %s, Size: %d bytes\n", model.Name, model.Size)
|
||||
}
|
||||
|
||||
// Pull a new model
|
||||
pullResp, err := sdk.PullModel(ctx, "codellama")
|
||||
fmt.Printf("Pull status: %s\n", pullResp.Status)
|
||||
|
||||
// Show model information
|
||||
modelInfo, err := sdk.ShowModel(ctx, "llama2")
|
||||
fmt.Printf("Model details: %+v\n", modelInfo.Details)
|
||||
|
||||
// Delete a model
|
||||
deleteResp, err := sdk.DeleteModel(ctx, "unwanted-model")
|
||||
fmt.Printf("Delete status: %s\n", deleteResp.Status)
|
||||
```
|
||||
|
||||
### Health Check
|
||||
|
||||
```go
|
||||
err := sdk.Health(ctx)
|
||||
if err != nil {
|
||||
log.Printf("Ollama service is not available: %v", err)
|
||||
} else {
|
||||
fmt.Println("Ollama service is healthy")
|
||||
}
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
The SDK provides specific error types for different failure scenarios:
|
||||
|
||||
```go
|
||||
resp, err := sdk.Generate(ctx, req)
|
||||
if err != nil {
|
||||
switch e := err.(type) {
|
||||
case *ollama.ConfigError:
|
||||
fmt.Printf("Configuration error: %s\n", e.Message)
|
||||
case *ollama.ModelError:
|
||||
fmt.Printf("Model error for %s: %s\n", e.Model, e.Message)
|
||||
case *ollama.GenerationError:
|
||||
fmt.Printf("Generation error: %s\n", e.Message)
|
||||
case *ollama.ConnectionError:
|
||||
fmt.Printf("Connection error: %s\n", e.Message)
|
||||
case *ollama.TimeoutError:
|
||||
fmt.Printf("Timeout during %s: %s\n", e.Operation, e.Message)
|
||||
default:
|
||||
fmt.Printf("Unexpected error: %v\n", err)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Custom HTTP Client
|
||||
|
||||
```go
|
||||
httpClient := &http.Client{
|
||||
Timeout: 60 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: 10,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
DisableCompression: true,
|
||||
},
|
||||
}
|
||||
|
||||
sdk, err := ollama.NewWithClient(config, httpClient, logger)
|
||||
```
|
||||
|
||||
### Custom Options
|
||||
|
||||
```go
|
||||
req := &ollama.GenerateRequest{
|
||||
Model: "llama2",
|
||||
Prompt: "Write a haiku",
|
||||
}
|
||||
|
||||
opts := &ollama.GenerateOptions{
|
||||
Temperature: &[]float64{0.5}[0],
|
||||
TopP: &[]float64{0.8}[0],
|
||||
TopK: &[]int{40}[0],
|
||||
RepeatPenalty: &[]float64{1.1}[0],
|
||||
Seed: &[]int{42}[0],
|
||||
Stop: []string{"</s>", "\n\n"},
|
||||
Format: "json",
|
||||
}
|
||||
|
||||
resp, err := sdk.GenerateWithOptions(ctx, req, opts)
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Always use context**: Pass context for cancellation and timeouts
|
||||
2. **Handle errors properly**: Use type assertions for specific error handling
|
||||
3. **Configure timeouts**: Set appropriate timeouts for your use case
|
||||
4. **Use streaming for long responses**: Enable streaming for better user experience
|
||||
5. **Monitor model usage**: Keep track of token usage and costs
|
||||
6. **Validate inputs**: Ensure prompts and parameters are within limits
|
||||
7. **Use structured logging**: Leverage the built-in logging for debugging
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Connection refused**: Ensure Ollama server is running on the correct port
|
||||
2. **Model not found**: Pull the model first using `PullModel()`
|
||||
3. **Timeout errors**: Increase the timeout configuration
|
||||
4. **Rate limiting**: Implement backoff strategies for high-volume usage
|
||||
|
||||
### Debug Mode
|
||||
|
||||
Enable debug logging to troubleshoot issues:
|
||||
|
||||
```go
|
||||
config := ollama.DefaultConfig()
|
||||
config.EnableLogging = true
|
||||
|
||||
loggerConfig := &logger.Config{
|
||||
Level: "debug",
|
||||
Format: "json",
|
||||
Output: "stdout",
|
||||
}
|
||||
projectLogger := logger.NewLogger(loggerConfig)
|
||||
|
||||
sdk, err := ollama.New(config, projectLogger)
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This SDK is part of the Tender Management system and follows the same licensing terms.
|
||||
@@ -0,0 +1,263 @@
|
||||
# Ollama SDK - Quick Usage Guide
|
||||
|
||||
## Installation
|
||||
The SDK is already part of your `tm` module. Import it as:
|
||||
```go
|
||||
import "tm/pkg/ollama"
|
||||
```
|
||||
|
||||
## Basic Setup
|
||||
```go
|
||||
// 1. Configuration
|
||||
config := ollama.DefaultConfig()
|
||||
config.BaseURL = "http://localhost:11434" // Your Ollama server URL
|
||||
config.DefaultModel = "llama2" // Default model to use
|
||||
|
||||
// 2. Logger (using project's logger)
|
||||
loggerConfig := &logger.Config{Level: "info", Format: "json", Output: "stdout"}
|
||||
projectLogger := logger.NewLogger(loggerConfig)
|
||||
|
||||
// 3. Create SDK
|
||||
sdk, err := ollama.New(config, projectLogger)
|
||||
if err != nil {
|
||||
log.Fatal("Failed to create Ollama SDK:", err)
|
||||
}
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
```bash
|
||||
export OLLAMA_BASE_URL="http://localhost:11434"
|
||||
export OLLAMA_TIMEOUT="300s"
|
||||
export OLLAMA_RETRY_ATTEMPTS="3"
|
||||
export OLLAMA_RETRY_DELAY="2s"
|
||||
export OLLAMA_ENABLE_LOGGING="true"
|
||||
export OLLAMA_DEFAULT_MODEL="llama2"
|
||||
export OLLAMA_DEFAULT_TEMPERATURE="0.7"
|
||||
export OLLAMA_DEFAULT_TOP_P="0.9"
|
||||
export OLLAMA_DEFAULT_MAX_TOKENS="1000"
|
||||
```
|
||||
|
||||
## Quick Examples
|
||||
|
||||
### Simple Text Generation
|
||||
```go
|
||||
ctx := context.Background()
|
||||
|
||||
// Generate text with default model
|
||||
response, err := sdk.QuickGenerate(ctx, "Write a short story about a robot")
|
||||
if err != nil {
|
||||
log.Printf("Generation failed: %v", err)
|
||||
}
|
||||
fmt.Println(response)
|
||||
|
||||
// Generate with specific model
|
||||
response, err := sdk.QuickGenerateWithModel(ctx, "codellama", "Write a Python function to sort a list")
|
||||
```
|
||||
|
||||
### Chat Conversations
|
||||
```go
|
||||
// Simple chat
|
||||
response, err := sdk.QuickChat(ctx, "What is the capital of France?")
|
||||
fmt.Println(response)
|
||||
|
||||
// Chat with system prompt
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a helpful coding assistant.",
|
||||
"How do I implement a binary search in Go?")
|
||||
```
|
||||
|
||||
### Advanced Chat with Builder
|
||||
```go
|
||||
// Create a chat builder
|
||||
builder := sdk.NewChatBuilder().
|
||||
SetModel("llama2").
|
||||
SetSystem("You are a helpful coding assistant.").
|
||||
SetTemperature(0.7).
|
||||
SetMaxTokens(500)
|
||||
|
||||
// Start conversation
|
||||
response, err := builder.
|
||||
AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?").
|
||||
Send(ctx)
|
||||
|
||||
fmt.Printf("Assistant: %s\n", response.Message.Content)
|
||||
|
||||
// Continue the conversation
|
||||
response, err = builder.Continue(ctx, "Can you show me an example with methods?")
|
||||
fmt.Printf("Assistant: %s\n", response.Message.Content)
|
||||
```
|
||||
|
||||
### Streaming Responses
|
||||
```go
|
||||
// Streaming text generation
|
||||
err := sdk.StreamGenerate(ctx, &ollama.GenerateRequest{
|
||||
Model: "llama2",
|
||||
Prompt: "Write a poem about programming",
|
||||
}, func(chunk string) error {
|
||||
fmt.Print(chunk) // Print each chunk as it arrives
|
||||
return nil
|
||||
})
|
||||
|
||||
// Streaming chat
|
||||
err := sdk.StreamChat(ctx, &ollama.ChatRequest{
|
||||
Model: "llama2",
|
||||
Messages: []ollama.Message{
|
||||
{
|
||||
Role: ollama.MessageRoleUser,
|
||||
Content: "Tell me a story",
|
||||
},
|
||||
},
|
||||
}, func(chunk string) error {
|
||||
fmt.Print(chunk)
|
||||
return nil
|
||||
})
|
||||
```
|
||||
|
||||
### Embeddings
|
||||
```go
|
||||
// Generate embeddings
|
||||
embeddings, err := sdk.QuickEmbed(ctx, "This is some text to embed")
|
||||
if err != nil {
|
||||
log.Printf("Embedding failed: %v", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Embedding dimension: %d\n", len(embeddings))
|
||||
```
|
||||
|
||||
### Model Management
|
||||
```go
|
||||
// List available models
|
||||
models, err := sdk.ListModels(ctx)
|
||||
for _, model := range models.Models {
|
||||
fmt.Printf("Model: %s, Size: %d bytes\n", model.Name, model.Size)
|
||||
}
|
||||
|
||||
// Pull a new model
|
||||
pullResp, err := sdk.PullModel(ctx, "codellama")
|
||||
fmt.Printf("Pull status: %s\n", pullResp.Status)
|
||||
|
||||
// Show model information
|
||||
modelInfo, err := sdk.ShowModel(ctx, "llama2")
|
||||
fmt.Printf("Model family: %s\n", modelInfo.Details.Family)
|
||||
|
||||
// Delete a model
|
||||
deleteResp, err := sdk.DeleteModel(ctx, "unwanted-model")
|
||||
fmt.Printf("Delete status: %s\n", deleteResp.Status)
|
||||
```
|
||||
|
||||
### Health Check
|
||||
```go
|
||||
err := sdk.Health(ctx)
|
||||
if err != nil {
|
||||
log.Printf("Ollama service is not available: %v", err)
|
||||
} else {
|
||||
fmt.Println("Ollama service is healthy")
|
||||
}
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
```go
|
||||
resp, err := sdk.Generate(ctx, req)
|
||||
if err != nil {
|
||||
switch e := err.(type) {
|
||||
case *ollama.ConfigError:
|
||||
fmt.Printf("Configuration error: %s\n", e.Message)
|
||||
case *ollama.ModelError:
|
||||
fmt.Printf("Model error for %s: %s\n", e.Model, e.Message)
|
||||
case *ollama.GenerationError:
|
||||
fmt.Printf("Generation error: %s\n", e.Message)
|
||||
case *ollama.ConnectionError:
|
||||
fmt.Printf("Connection error: %s\n", e.Message)
|
||||
case *ollama.TimeoutError:
|
||||
fmt.Printf("Timeout during %s: %s\n", e.Operation, e.Message)
|
||||
default:
|
||||
fmt.Printf("Unexpected error: %v\n", err)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Advanced Configuration
|
||||
```go
|
||||
config := &ollama.Config{
|
||||
BaseURL: "http://localhost:11434",
|
||||
Timeout: 300 * time.Second,
|
||||
RetryAttempts: 3,
|
||||
RetryDelay: 2 * time.Second,
|
||||
EnableLogging: true,
|
||||
UserAgent: "MyApp-OllamaSDK/1.0",
|
||||
DefaultModel: "llama2",
|
||||
DefaultTemperature: 0.7,
|
||||
DefaultTopP: 0.9,
|
||||
DefaultMaxTokens: 1000,
|
||||
EnableStreaming: false,
|
||||
}
|
||||
|
||||
sdk, err := ollama.New(config, logger)
|
||||
```
|
||||
|
||||
## Custom HTTP Client
|
||||
```go
|
||||
httpClient := &http.Client{
|
||||
Timeout: 60 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: 10,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
DisableCompression: true,
|
||||
},
|
||||
}
|
||||
|
||||
sdk, err := ollama.NewWithClient(config, httpClient, logger)
|
||||
```
|
||||
|
||||
## Integration with Project Logger
|
||||
```go
|
||||
// Use the project's logger system
|
||||
loggerConfig := &logger.Config{
|
||||
Level: "info",
|
||||
Format: "json",
|
||||
Output: "stdout",
|
||||
}
|
||||
projectLogger := logger.NewLogger(loggerConfig)
|
||||
|
||||
sdk, err := ollama.New(config, projectLogger)
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Always use context**: Pass context for cancellation and timeouts
|
||||
2. **Handle errors properly**: Use type assertions for specific error handling
|
||||
3. **Configure timeouts**: Set appropriate timeouts for your use case
|
||||
4. **Use streaming for long responses**: Enable streaming for better user experience
|
||||
5. **Monitor model usage**: Keep track of token usage and costs
|
||||
6. **Validate inputs**: Ensure prompts and parameters are within limits
|
||||
7. **Use structured logging**: Leverage the built-in logging for debugging
|
||||
|
||||
## Common Use Cases
|
||||
|
||||
### Document Summarization
|
||||
```go
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a helpful assistant that summarizes documents concisely.",
|
||||
"Please summarize this document: [document content]")
|
||||
```
|
||||
|
||||
### Code Review Assistant
|
||||
```go
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a senior software engineer reviewing code. Provide constructive feedback.",
|
||||
"Please review this Go code: [code snippet]")
|
||||
```
|
||||
|
||||
### Translation Service
|
||||
```go
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a professional translator. Translate the following text accurately.",
|
||||
"Translate this to French: Hello, how are you?")
|
||||
```
|
||||
|
||||
### Question Answering
|
||||
```go
|
||||
response, err := sdk.QuickChatWithSystem(ctx,
|
||||
"You are a helpful assistant that answers questions accurately and concisely.",
|
||||
"What is the difference between HTTP and HTTPS?")
|
||||
```
|
||||
@@ -0,0 +1,290 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
"tm/pkg/logger"
|
||||
)
|
||||
|
||||
// Client implements the HTTPClient interface for Ollama API
|
||||
type Client struct {
|
||||
config *Config
|
||||
httpClient *http.Client
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
// NewClient creates a new Ollama HTTP client
|
||||
func NewClient(config *Config, logger logger.Logger) (*Client, error) {
|
||||
if config == nil {
|
||||
config = DefaultConfig()
|
||||
}
|
||||
|
||||
if err := config.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
httpClient := &http.Client{
|
||||
Timeout: config.Timeout,
|
||||
}
|
||||
|
||||
return &Client{
|
||||
config: config,
|
||||
httpClient: httpClient,
|
||||
logger: logger,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NewClientWithHTTPClient creates a new client with a custom HTTP client
|
||||
func NewClientWithHTTPClient(config *Config, httpClient *http.Client, logger logger.Logger) (*Client, error) {
|
||||
if config == nil {
|
||||
config = DefaultConfig()
|
||||
}
|
||||
|
||||
if err := config.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Client{
|
||||
config: config,
|
||||
httpClient: httpClient,
|
||||
logger: logger,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Post performs a POST request to the Ollama API
|
||||
func (c *Client) Post(ctx context.Context, url string, body interface{}, result interface{}) error {
|
||||
return c.makeRequest(ctx, "POST", url, body, result)
|
||||
}
|
||||
|
||||
// Get performs a GET request to the Ollama API
|
||||
func (c *Client) Get(ctx context.Context, url string, result interface{}) error {
|
||||
return c.makeRequest(ctx, "GET", url, nil, result)
|
||||
}
|
||||
|
||||
// Delete performs a DELETE request to the Ollama API
|
||||
func (c *Client) Delete(ctx context.Context, url string, body interface{}, result interface{}) error {
|
||||
return c.makeRequest(ctx, "DELETE", url, body, result)
|
||||
}
|
||||
|
||||
// makeRequest performs an HTTP request with retry logic and logging
|
||||
func (c *Client) makeRequest(ctx context.Context, method, endpoint string, body interface{}, result interface{}) error {
|
||||
url := c.config.BaseURL + endpoint
|
||||
|
||||
var bodyReader io.Reader
|
||||
if body != nil {
|
||||
bodyBytes, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
bodyReader = bytes.NewReader(bodyBytes)
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for attempt := 0; attempt <= c.config.RetryAttempts; attempt++ {
|
||||
if attempt > 0 {
|
||||
// Wait before retry
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(c.config.RetryDelay):
|
||||
}
|
||||
}
|
||||
|
||||
err := c.doRequest(ctx, method, url, bodyReader, result)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
lastErr = err
|
||||
|
||||
// Don't retry on certain errors
|
||||
if c.shouldNotRetry(err) {
|
||||
break
|
||||
}
|
||||
|
||||
if c.logger != nil {
|
||||
c.logger.Warn("Request failed, retrying", map[string]interface{}{
|
||||
"attempt": attempt + 1,
|
||||
"max_attempts": c.config.RetryAttempts + 1,
|
||||
"error": err.Error(),
|
||||
"url": url,
|
||||
"method": method,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Errorf("request failed after %d attempts: %w", c.config.RetryAttempts+1, lastErr)
|
||||
}
|
||||
|
||||
// doRequest performs a single HTTP request
|
||||
func (c *Client) doRequest(ctx context.Context, method, url string, body io.Reader, result interface{}) error {
|
||||
req, err := http.NewRequestWithContext(ctx, method, url, body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", c.config.UserAgent)
|
||||
|
||||
if c.logger != nil && c.config.EnableLogging {
|
||||
c.logger.Debug("Making HTTP request", map[string]interface{}{
|
||||
"method": method,
|
||||
"url": url,
|
||||
"user_agent": c.config.UserAgent,
|
||||
})
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if c.logger != nil && c.config.EnableLogging {
|
||||
c.logger.Debug("Received HTTP response", map[string]interface{}{
|
||||
"status_code": resp.StatusCode,
|
||||
"url": url,
|
||||
})
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
if result != nil {
|
||||
if err := json.NewDecoder(resp.Body).Decode(result); err != nil {
|
||||
return fmt.Errorf("failed to decode response: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// shouldNotRetry determines if an error should not be retried
|
||||
func (c *Client) shouldNotRetry(err error) bool {
|
||||
// Don't retry on context cancellation
|
||||
if err == context.Canceled || err == context.DeadlineExceeded {
|
||||
return true
|
||||
}
|
||||
|
||||
// Don't retry on certain HTTP status codes
|
||||
if err.Error() == "HTTP 400" || err.Error() == "HTTP 401" || err.Error() == "HTTP 403" || err.Error() == "HTTP 404" {
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// StreamGenerate performs a streaming text generation request
|
||||
func (c *Client) StreamGenerate(ctx context.Context, req *GenerateRequest, onChunk func(*StreamChunk) error) error {
|
||||
url := c.config.BaseURL + "/api/generate"
|
||||
|
||||
// Force streaming
|
||||
stream := true
|
||||
req.Stream = &stream
|
||||
|
||||
bodyBytes, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", c.config.UserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
for {
|
||||
var chunk StreamChunk
|
||||
if err := decoder.Decode(&chunk); err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return fmt.Errorf("failed to decode stream chunk: %w", err)
|
||||
}
|
||||
|
||||
if err := onChunk(&chunk); err != nil {
|
||||
return fmt.Errorf("error processing stream chunk: %w", err)
|
||||
}
|
||||
|
||||
if chunk.Done {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// StreamChat performs a streaming chat request
|
||||
func (c *Client) StreamChat(ctx context.Context, req *ChatRequest, onChunk func(*StreamChunk) error) error {
|
||||
url := c.config.BaseURL + "/api/chat"
|
||||
|
||||
// Force streaming
|
||||
stream := true
|
||||
req.Stream = &stream
|
||||
|
||||
bodyBytes, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", c.config.UserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to execute request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
for {
|
||||
var chunk StreamChunk
|
||||
if err := decoder.Decode(&chunk); err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
return fmt.Errorf("failed to decode stream chunk: %w", err)
|
||||
}
|
||||
|
||||
if err := onChunk(&chunk); err != nil {
|
||||
return fmt.Errorf("error processing stream chunk: %w", err)
|
||||
}
|
||||
|
||||
if chunk.Done {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config represents the configuration for the Ollama SDK
|
||||
type Config struct {
|
||||
// BaseURL is the base URL of the Ollama API server
|
||||
BaseURL string `env:"OLLAMA_BASE_URL" default:"http://localhost:11434"`
|
||||
|
||||
// Timeout is the timeout for HTTP requests
|
||||
Timeout time.Duration `env:"OLLAMA_TIMEOUT" default:"300s"`
|
||||
|
||||
// RetryAttempts is the number of retry attempts for failed requests
|
||||
RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS" default:"3"`
|
||||
|
||||
// RetryDelay is the delay between retry attempts
|
||||
RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY" default:"2s"`
|
||||
|
||||
// EnableLogging enables request/response logging
|
||||
EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING" default:"true"`
|
||||
|
||||
// UserAgent is the user agent string for HTTP requests
|
||||
UserAgent string `env:"OLLAMA_USER_AGENT" default:"TenderManagement-OllamaSDK/1.0"`
|
||||
|
||||
// DefaultModel is the default model to use if none is specified
|
||||
DefaultModel string `env:"OLLAMA_DEFAULT_MODEL" default:"llama2"`
|
||||
|
||||
// DefaultTemperature is the default temperature for text generation
|
||||
DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE" default:"0.7"`
|
||||
|
||||
// DefaultTopP is the default top_p value for text generation
|
||||
DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P" default:"0.9"`
|
||||
|
||||
// DefaultMaxTokens is the default maximum number of tokens to generate
|
||||
DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS" default:"1000"`
|
||||
|
||||
// EnableStreaming enables streaming responses by default
|
||||
EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING" default:"false"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns a default configuration for the Ollama SDK
|
||||
func DefaultConfig() *Config {
|
||||
return &Config{
|
||||
BaseURL: "http://localhost:11434",
|
||||
Timeout: 300 * time.Second,
|
||||
RetryAttempts: 3,
|
||||
RetryDelay: 2 * time.Second,
|
||||
EnableLogging: true,
|
||||
UserAgent: "TenderManagement-OllamaSDK/1.0",
|
||||
DefaultModel: "llama2",
|
||||
DefaultTemperature: 0.7,
|
||||
DefaultTopP: 0.9,
|
||||
DefaultMaxTokens: 1000,
|
||||
EnableStreaming: false,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate validates the configuration
|
||||
func (c *Config) Validate() error {
|
||||
if c.BaseURL == "" {
|
||||
return ErrInvalidConfig("BaseURL cannot be empty")
|
||||
}
|
||||
|
||||
if c.Timeout <= 0 {
|
||||
return ErrInvalidConfig("Timeout must be positive")
|
||||
}
|
||||
|
||||
if c.RetryAttempts < 0 {
|
||||
return ErrInvalidConfig("RetryAttempts cannot be negative")
|
||||
}
|
||||
|
||||
if c.RetryDelay < 0 {
|
||||
return ErrInvalidConfig("RetryDelay cannot be negative")
|
||||
}
|
||||
|
||||
if c.DefaultTemperature < 0 || c.DefaultTemperature > 2 {
|
||||
return ErrInvalidConfig("DefaultTemperature must be between 0 and 2")
|
||||
}
|
||||
|
||||
if c.DefaultTopP < 0 || c.DefaultTopP > 1 {
|
||||
return ErrInvalidConfig("DefaultTopP must be between 0 and 1")
|
||||
}
|
||||
|
||||
if c.DefaultMaxTokens < 0 {
|
||||
return ErrInvalidConfig("DefaultMaxTokens cannot be negative")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// GenerateRequest represents a request to generate text
|
||||
type GenerateRequest struct {
|
||||
Model string `json:"model" validate:"required"`
|
||||
Prompt string `json:"prompt" validate:"required"`
|
||||
Stream *bool `json:"stream,omitempty"`
|
||||
Format string `json:"format,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
System string `json:"system,omitempty"`
|
||||
Template string `json:"template,omitempty"`
|
||||
Context []int `json:"context,omitempty"`
|
||||
Raw bool `json:"raw,omitempty"`
|
||||
KeepAlive string `json:"keep_alive,omitempty"`
|
||||
}
|
||||
|
||||
// GenerateResponse represents the response from a text generation request
|
||||
type GenerateResponse struct {
|
||||
Model string `json:"model"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Response string `json:"response"`
|
||||
Done bool `json:"done"`
|
||||
Context []int `json:"context,omitempty"`
|
||||
TotalDuration int64 `json:"total_duration,omitempty"`
|
||||
LoadDuration int64 `json:"load_duration,omitempty"`
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
|
||||
PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"`
|
||||
EvalCount int `json:"eval_count,omitempty"`
|
||||
EvalDuration int64 `json:"eval_duration,omitempty"`
|
||||
}
|
||||
|
||||
// ChatRequest represents a request for a chat conversation
|
||||
type ChatRequest struct {
|
||||
Model string `json:"model" validate:"required"`
|
||||
Messages []Message `json:"messages" validate:"required,min=1"`
|
||||
Stream *bool `json:"stream,omitempty"`
|
||||
Format string `json:"format,omitempty"`
|
||||
Options ChatOptions `json:"options,omitempty"`
|
||||
KeepAlive string `json:"keep_alive,omitempty"`
|
||||
}
|
||||
|
||||
// Message represents a message in a chat conversation
|
||||
type Message struct {
|
||||
Role MessageRole `json:"role" validate:"required,oneof=system user assistant"`
|
||||
Content string `json:"content" validate:"required"`
|
||||
Images []string `json:"images,omitempty"`
|
||||
}
|
||||
|
||||
// ChatOptions represents options for chat requests
|
||||
type ChatOptions struct {
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
RepeatPenalty *float64 `json:"repeat_penalty,omitempty"`
|
||||
Seed *int `json:"seed,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
NumCtx *int `json:"num_ctx,omitempty"`
|
||||
NumPredict *int `json:"num_predict,omitempty"`
|
||||
NumKeep *int `json:"num_keep,omitempty"`
|
||||
NumThread *int `json:"num_thread,omitempty"`
|
||||
RepeatLastN *int `json:"repeat_last_n,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// ChatResponse represents the response from a chat request
|
||||
type ChatResponse struct {
|
||||
Model string `json:"model"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Message Message `json:"message"`
|
||||
Done bool `json:"done"`
|
||||
TotalDuration int64 `json:"total_duration,omitempty"`
|
||||
LoadDuration int64 `json:"load_duration,omitempty"`
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
|
||||
PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"`
|
||||
EvalCount int `json:"eval_count,omitempty"`
|
||||
EvalDuration int64 `json:"eval_duration,omitempty"`
|
||||
}
|
||||
|
||||
// EmbedRequest represents a request to generate embeddings
|
||||
type EmbedRequest struct {
|
||||
Model string `json:"model" validate:"required"`
|
||||
Prompt string `json:"prompt" validate:"required"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// EmbedResponse represents the response from an embedding request
|
||||
type EmbedResponse struct {
|
||||
Embedding []float64 `json:"embedding"`
|
||||
}
|
||||
|
||||
// PullRequest represents a request to pull a model
|
||||
type PullRequest struct {
|
||||
Name string `json:"name" validate:"required"`
|
||||
Insecure bool `json:"insecure,omitempty"`
|
||||
Stream *bool `json:"stream,omitempty"`
|
||||
}
|
||||
|
||||
// PullResponse represents the response from a pull request
|
||||
type PullResponse struct {
|
||||
Status string `json:"status"`
|
||||
Digest string `json:"digest,omitempty"`
|
||||
Total int64 `json:"total,omitempty"`
|
||||
Completed int64 `json:"completed,omitempty"`
|
||||
}
|
||||
|
||||
// ListModelsResponse represents the response from listing models
|
||||
type ListModelsResponse struct {
|
||||
Models []ModelInfo `json:"models"`
|
||||
}
|
||||
|
||||
// ModelInfo represents information about a model
|
||||
type ModelInfo struct {
|
||||
Name string `json:"name"`
|
||||
ModifiedAt time.Time `json:"modified_at"`
|
||||
Size int64 `json:"size"`
|
||||
Digest string `json:"digest"`
|
||||
Details ModelDetails `json:"details,omitempty"`
|
||||
}
|
||||
|
||||
// ModelDetails represents detailed information about a model
|
||||
type ModelDetails struct {
|
||||
Format string `json:"format,omitempty"`
|
||||
Family string `json:"family,omitempty"`
|
||||
Families []string `json:"families,omitempty"`
|
||||
ParameterSize string `json:"parameter_size,omitempty"`
|
||||
QuantizationLevel string `json:"quantization_level,omitempty"`
|
||||
ParentModel string `json:"parent_model,omitempty"`
|
||||
Templates map[string]interface{} `json:"templates,omitempty"`
|
||||
System string `json:"system,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// DeleteRequest represents a request to delete a model
|
||||
type DeleteRequest struct {
|
||||
Name string `json:"name" validate:"required"`
|
||||
}
|
||||
|
||||
// DeleteResponse represents the response from a delete request
|
||||
type DeleteResponse struct {
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// ShowModelRequest represents a request to show model information
|
||||
type ShowModelRequest struct {
|
||||
Name string `json:"name" validate:"required"`
|
||||
Verbose bool `json:"verbose,omitempty"`
|
||||
}
|
||||
|
||||
// ShowModelResponse represents the response from a show model request
|
||||
type ShowModelResponse struct {
|
||||
License string `json:"license,omitempty"`
|
||||
Modelfile string `json:"modelfile,omitempty"`
|
||||
Parameters string `json:"parameters,omitempty"`
|
||||
Template string `json:"template,omitempty"`
|
||||
Details ModelDetails `json:"details,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// HealthResponse represents the health check response
|
||||
type HealthResponse struct {
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// StreamChunk represents a chunk in a streaming response
|
||||
type StreamChunk struct {
|
||||
Model string `json:"model,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at,omitempty"`
|
||||
Response string `json:"response,omitempty"`
|
||||
Done bool `json:"done"`
|
||||
Context []int `json:"context,omitempty"`
|
||||
TotalDuration int64 `json:"total_duration,omitempty"`
|
||||
LoadDuration int64 `json:"load_duration,omitempty"`
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
|
||||
PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"`
|
||||
EvalCount int `json:"eval_count,omitempty"`
|
||||
EvalDuration int64 `json:"eval_duration,omitempty"`
|
||||
Message *Message `json:"message,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Error types for the Ollama SDK
|
||||
var (
|
||||
// ErrInvalidConfig is returned when configuration is invalid
|
||||
ErrInvalidConfig = func(msg string) error {
|
||||
return &ConfigError{Message: msg}
|
||||
}
|
||||
|
||||
// ErrInvalidRequest is returned when a request is invalid
|
||||
ErrInvalidRequest = func(msg string) error {
|
||||
return &RequestError{Message: msg}
|
||||
}
|
||||
|
||||
// ErrModelNotFound is returned when a model is not found
|
||||
ErrModelNotFound = func(model string) error {
|
||||
return &ModelError{Message: fmt.Sprintf("model %s not found", model), Model: model}
|
||||
}
|
||||
|
||||
// ErrModelPullFailed is returned when pulling a model fails
|
||||
ErrModelPullFailed = func(model string, reason string) error {
|
||||
return &ModelError{Message: fmt.Sprintf("failed to pull model %s: %s", model, reason), Model: model}
|
||||
}
|
||||
|
||||
// ErrModelDeleteFailed is returned when deleting a model fails
|
||||
ErrModelDeleteFailed = func(model string, reason string) error {
|
||||
return &ModelError{Message: fmt.Sprintf("failed to delete model %s: %s", model, reason), Model: model}
|
||||
}
|
||||
|
||||
// ErrGenerationFailed is returned when text generation fails
|
||||
ErrGenerationFailed = func(reason string) error {
|
||||
return &GenerationError{Message: fmt.Sprintf("generation failed: %s", reason)}
|
||||
}
|
||||
|
||||
// ErrChatFailed is returned when chat fails
|
||||
ErrChatFailed = func(reason string) error {
|
||||
return &ChatError{Message: fmt.Sprintf("chat failed: %s", reason)}
|
||||
}
|
||||
|
||||
// ErrEmbeddingFailed is returned when embedding generation fails
|
||||
ErrEmbeddingFailed = func(reason string) error {
|
||||
return &EmbeddingError{Message: fmt.Sprintf("embedding failed: %s", reason)}
|
||||
}
|
||||
|
||||
// ErrConnectionFailed is returned when connection to Ollama fails
|
||||
ErrConnectionFailed = func(reason string) error {
|
||||
return &ConnectionError{Message: fmt.Sprintf("connection failed: %s", reason)}
|
||||
}
|
||||
|
||||
// ErrTimeout is returned when a request times out
|
||||
ErrTimeout = func(operation string) error {
|
||||
return &TimeoutError{Message: fmt.Sprintf("timeout during %s", operation), Operation: operation}
|
||||
}
|
||||
|
||||
// ErrRateLimitExceeded is returned when rate limit is exceeded
|
||||
ErrRateLimitExceeded = func() error {
|
||||
return &RateLimitError{Message: "rate limit exceeded"}
|
||||
}
|
||||
)
|
||||
|
||||
// ConfigError represents a configuration error
|
||||
type ConfigError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *ConfigError) Error() string {
|
||||
return fmt.Sprintf("config error: %s", e.Message)
|
||||
}
|
||||
|
||||
// RequestError represents a request validation error
|
||||
type RequestError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *RequestError) Error() string {
|
||||
return fmt.Sprintf("request error: %s", e.Message)
|
||||
}
|
||||
|
||||
// ModelError represents a model-related error
|
||||
type ModelError struct {
|
||||
Message string
|
||||
Model string
|
||||
}
|
||||
|
||||
func (e *ModelError) Error() string {
|
||||
return fmt.Sprintf("model error: %s", e.Message)
|
||||
}
|
||||
|
||||
// GenerationError represents a text generation error
|
||||
type GenerationError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *GenerationError) Error() string {
|
||||
return fmt.Sprintf("generation error: %s", e.Message)
|
||||
}
|
||||
|
||||
// ChatError represents a chat error
|
||||
type ChatError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *ChatError) Error() string {
|
||||
return fmt.Sprintf("chat error: %s", e.Message)
|
||||
}
|
||||
|
||||
// EmbeddingError represents an embedding generation error
|
||||
type EmbeddingError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *EmbeddingError) Error() string {
|
||||
return fmt.Sprintf("embedding error: %s", e.Message)
|
||||
}
|
||||
|
||||
// ConnectionError represents a connection error
|
||||
type ConnectionError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *ConnectionError) Error() string {
|
||||
return fmt.Sprintf("connection error: %s", e.Message)
|
||||
}
|
||||
|
||||
// TimeoutError represents a timeout error
|
||||
type TimeoutError struct {
|
||||
Message string
|
||||
Operation string
|
||||
}
|
||||
|
||||
func (e *TimeoutError) Error() string {
|
||||
return fmt.Sprintf("timeout error: %s", e.Message)
|
||||
}
|
||||
|
||||
// RateLimitError represents a rate limit error
|
||||
type RateLimitError struct {
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *RateLimitError) Error() string {
|
||||
return fmt.Sprintf("rate limit error: %s", e.Message)
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
// OllamaService defines the interface for Ollama operations
|
||||
type OllamaService interface {
|
||||
// Generate generates text using a specified model and prompt
|
||||
Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error)
|
||||
|
||||
// GenerateWithOptions generates text with additional options
|
||||
GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error)
|
||||
|
||||
// Chat performs a conversational chat with the model
|
||||
Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error)
|
||||
|
||||
// Embed generates embeddings for the given text
|
||||
Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error)
|
||||
|
||||
// PullModel pulls a model from the registry
|
||||
PullModel(ctx context.Context, model string) (*PullResponse, error)
|
||||
|
||||
// ListModels lists available models
|
||||
ListModels(ctx context.Context) (*ListModelsResponse, error)
|
||||
|
||||
// DeleteModel deletes a model
|
||||
DeleteModel(ctx context.Context, model string) (*DeleteResponse, error)
|
||||
|
||||
// ShowModel shows model information
|
||||
ShowModel(ctx context.Context, model string) (*ShowModelResponse, error)
|
||||
|
||||
// Health checks if the Ollama service is available
|
||||
Health(ctx context.Context) error
|
||||
|
||||
// NewChatBuilder creates a new chat builder for conversational interactions
|
||||
NewChatBuilder() ChatBuilder
|
||||
}
|
||||
|
||||
// HTTPClient defines the interface for HTTP operations
|
||||
type HTTPClient interface {
|
||||
// Post performs a POST request
|
||||
Post(ctx context.Context, url string, body interface{}, result interface{}) error
|
||||
|
||||
// Get performs a GET request
|
||||
Get(ctx context.Context, url string, result interface{}) error
|
||||
|
||||
// Delete performs a DELETE request
|
||||
Delete(ctx context.Context, url string, result interface{}) error
|
||||
}
|
||||
|
||||
// Logger defines the interface for logging operations
|
||||
type Logger interface {
|
||||
Debug(msg string, fields map[string]interface{})
|
||||
Info(msg string, fields map[string]interface{})
|
||||
Warn(msg string, fields map[string]interface{})
|
||||
Error(msg string, fields map[string]interface{})
|
||||
Fatal(msg string, fields map[string]interface{})
|
||||
WithFields(fields map[string]interface{}) Logger
|
||||
Sync() error
|
||||
}
|
||||
|
||||
// ChatBuilder provides a fluent interface for building chat requests
|
||||
type ChatBuilder interface {
|
||||
// SetModel sets the model to use
|
||||
SetModel(model string) ChatBuilder
|
||||
|
||||
// SetSystem sets the system prompt
|
||||
SetSystem(system string) ChatBuilder
|
||||
|
||||
// AddMessage adds a message to the conversation
|
||||
AddMessage(role MessageRole, content string) ChatBuilder
|
||||
|
||||
// SetTemperature sets the temperature for generation
|
||||
SetTemperature(temperature float64) ChatBuilder
|
||||
|
||||
// SetTopP sets the top_p parameter
|
||||
SetTopP(topP float64) ChatBuilder
|
||||
|
||||
// SetMaxTokens sets the maximum number of tokens to generate
|
||||
SetMaxTokens(maxTokens int) ChatBuilder
|
||||
|
||||
// SetStream sets whether to stream the response
|
||||
SetStream(stream bool) ChatBuilder
|
||||
|
||||
// SetFormat sets the response format
|
||||
SetFormat(format string) ChatBuilder
|
||||
|
||||
// SetOptions sets custom options
|
||||
SetOptions(options map[string]interface{}) ChatBuilder
|
||||
|
||||
// Build creates the chat request
|
||||
Build() (*ChatRequest, error)
|
||||
|
||||
// Send builds and sends the chat request
|
||||
Send(ctx context.Context) (*ChatResponse, error)
|
||||
|
||||
// Continue continues the conversation with the current context
|
||||
Continue(ctx context.Context, userMessage string) (*ChatResponse, error)
|
||||
}
|
||||
|
||||
// MessageRole represents the role of a message in a chat
|
||||
type MessageRole string
|
||||
|
||||
const (
|
||||
MessageRoleSystem MessageRole = "system"
|
||||
MessageRoleUser MessageRole = "user"
|
||||
MessageRoleAssistant MessageRole = "assistant"
|
||||
)
|
||||
|
||||
// GenerateOptions represents additional options for text generation
|
||||
type GenerateOptions struct {
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
RepeatPenalty *float64 `json:"repeat_penalty,omitempty"`
|
||||
Seed *int `json:"seed,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
Format string `json:"format,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"tm/pkg/logger"
|
||||
)
|
||||
|
||||
// SDK provides a high-level interface for the Ollama service
|
||||
type SDK struct {
|
||||
service OllamaService
|
||||
client *Client
|
||||
config *Config
|
||||
}
|
||||
|
||||
// New creates a new Ollama SDK with the provided configuration and logger
|
||||
func New(config *Config, logger logger.Logger) (*SDK, error) {
|
||||
if config == nil {
|
||||
config = DefaultConfig()
|
||||
}
|
||||
|
||||
client, err := NewClient(config, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create ollama client: %w", err)
|
||||
}
|
||||
|
||||
service := NewService(client)
|
||||
|
||||
return &SDK{
|
||||
service: service,
|
||||
client: client,
|
||||
config: config,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NewWithClient creates a new Ollama SDK with a custom HTTP client
|
||||
func NewWithClient(config *Config, httpClient *http.Client, logger logger.Logger) (*SDK, error) {
|
||||
if config == nil {
|
||||
config = DefaultConfig()
|
||||
}
|
||||
|
||||
client, err := NewClientWithHTTPClient(config, httpClient, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create ollama client: %w", err)
|
||||
}
|
||||
|
||||
service := NewService(client)
|
||||
|
||||
return &SDK{
|
||||
service: service,
|
||||
client: client,
|
||||
config: config,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Generate generates text using a specified model and prompt
|
||||
func (s *SDK) Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error) {
|
||||
return s.service.Generate(ctx, req)
|
||||
}
|
||||
|
||||
// GenerateWithOptions generates text with additional options
|
||||
func (s *SDK) GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error) {
|
||||
return s.service.GenerateWithOptions(ctx, req, opts)
|
||||
}
|
||||
|
||||
// Chat performs a conversational chat with the model
|
||||
func (s *SDK) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
|
||||
return s.service.Chat(ctx, req)
|
||||
}
|
||||
|
||||
// Embed generates embeddings for the given text
|
||||
func (s *SDK) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) {
|
||||
return s.service.Embed(ctx, req)
|
||||
}
|
||||
|
||||
// PullModel pulls a model from the registry
|
||||
func (s *SDK) PullModel(ctx context.Context, model string) (*PullResponse, error) {
|
||||
return s.service.PullModel(ctx, model)
|
||||
}
|
||||
|
||||
// ListModels lists available models
|
||||
func (s *SDK) ListModels(ctx context.Context) (*ListModelsResponse, error) {
|
||||
return s.service.ListModels(ctx)
|
||||
}
|
||||
|
||||
// DeleteModel deletes a model
|
||||
func (s *SDK) DeleteModel(ctx context.Context, model string) (*DeleteResponse, error) {
|
||||
return s.service.DeleteModel(ctx, model)
|
||||
}
|
||||
|
||||
// ShowModel shows model information
|
||||
func (s *SDK) ShowModel(ctx context.Context, model string) (*ShowModelResponse, error) {
|
||||
return s.service.ShowModel(ctx, model)
|
||||
}
|
||||
|
||||
// Health checks if the Ollama service is available
|
||||
func (s *SDK) Health(ctx context.Context) error {
|
||||
return s.service.Health(ctx)
|
||||
}
|
||||
|
||||
// NewChatBuilder creates a new chat builder for conversational interactions
|
||||
func (s *SDK) NewChatBuilder() ChatBuilder {
|
||||
return s.service.NewChatBuilder()
|
||||
}
|
||||
|
||||
// QuickGenerate is a convenience method for simple text generation
|
||||
func (s *SDK) QuickGenerate(ctx context.Context, prompt string) (string, error) {
|
||||
req := &GenerateRequest{
|
||||
Model: s.config.DefaultModel,
|
||||
Prompt: prompt,
|
||||
}
|
||||
|
||||
resp, err := s.Generate(ctx, req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return resp.Response, nil
|
||||
}
|
||||
|
||||
// QuickGenerateWithModel is a convenience method for text generation with a specific model
|
||||
func (s *SDK) QuickGenerateWithModel(ctx context.Context, model, prompt string) (string, error) {
|
||||
req := &GenerateRequest{
|
||||
Model: model,
|
||||
Prompt: prompt,
|
||||
}
|
||||
|
||||
resp, err := s.Generate(ctx, req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return resp.Response, nil
|
||||
}
|
||||
|
||||
// QuickChat is a convenience method for simple chat interactions
|
||||
func (s *SDK) QuickChat(ctx context.Context, userMessage string) (string, error) {
|
||||
req := &ChatRequest{
|
||||
Model: s.config.DefaultModel,
|
||||
Messages: []Message{
|
||||
{
|
||||
Role: MessageRoleUser,
|
||||
Content: userMessage,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := s.Chat(ctx, req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return resp.Message.Content, nil
|
||||
}
|
||||
|
||||
// QuickChatWithSystem is a convenience method for chat with system prompt
|
||||
func (s *SDK) QuickChatWithSystem(ctx context.Context, systemPrompt, userMessage string) (string, error) {
|
||||
req := &ChatRequest{
|
||||
Model: s.config.DefaultModel,
|
||||
Messages: []Message{
|
||||
{
|
||||
Role: MessageRoleSystem,
|
||||
Content: systemPrompt,
|
||||
},
|
||||
{
|
||||
Role: MessageRoleUser,
|
||||
Content: userMessage,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := s.Chat(ctx, req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return resp.Message.Content, nil
|
||||
}
|
||||
|
||||
// QuickEmbed is a convenience method for generating embeddings
|
||||
func (s *SDK) QuickEmbed(ctx context.Context, text string) ([]float64, error) {
|
||||
req := &EmbedRequest{
|
||||
Model: s.config.DefaultModel,
|
||||
Prompt: text,
|
||||
}
|
||||
|
||||
resp, err := s.Embed(ctx, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return resp.Embedding, nil
|
||||
}
|
||||
|
||||
// StreamGenerate performs streaming text generation
|
||||
func (s *SDK) StreamGenerate(ctx context.Context, req *GenerateRequest, onChunk func(string) error) error {
|
||||
// Force streaming
|
||||
stream := true
|
||||
req.Stream = &stream
|
||||
|
||||
return s.client.StreamGenerate(ctx, req, func(chunk *StreamChunk) error {
|
||||
if chunk.Response != "" {
|
||||
return onChunk(chunk.Response)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// StreamChat performs streaming chat
|
||||
func (s *SDK) StreamChat(ctx context.Context, req *ChatRequest, onChunk func(string) error) error {
|
||||
// Force streaming
|
||||
stream := true
|
||||
req.Stream = &stream
|
||||
|
||||
return s.client.StreamChat(ctx, req, func(chunk *StreamChunk) error {
|
||||
if chunk.Message != nil && chunk.Message.Content != "" {
|
||||
return onChunk(chunk.Message.Content)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// GetService returns the underlying service for advanced usage
|
||||
func (s *SDK) GetService() OllamaService {
|
||||
return s.service
|
||||
}
|
||||
|
||||
// GetConfig returns the current configuration
|
||||
func (s *SDK) GetConfig() *Config {
|
||||
return s.config
|
||||
}
|
||||
@@ -0,0 +1,444 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"tm/pkg/logger"
|
||||
)
|
||||
|
||||
// Service implements the OllamaService interface
|
||||
type Service struct {
|
||||
client *Client
|
||||
config *Config
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
// NewService creates a new Ollama service
|
||||
func NewService(client *Client) OllamaService {
|
||||
return &Service{
|
||||
client: client,
|
||||
config: client.config,
|
||||
logger: client.logger,
|
||||
}
|
||||
}
|
||||
|
||||
// Generate generates text using a specified model and prompt
|
||||
func (s *Service) Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Generating text", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"prompt_length": len(req.Prompt),
|
||||
})
|
||||
}
|
||||
|
||||
// Set defaults
|
||||
if req.Model == "" {
|
||||
req.Model = s.config.DefaultModel
|
||||
}
|
||||
|
||||
var resp GenerateResponse
|
||||
err := s.client.Post(ctx, "/api/generate", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Text generation failed", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, ErrGenerationFailed(err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Text generation completed", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"response_length": len(resp.Response),
|
||||
"eval_count": resp.EvalCount,
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// GenerateWithOptions generates text with additional options
|
||||
func (s *Service) GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Generating text with options", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"prompt_length": len(req.Prompt),
|
||||
})
|
||||
}
|
||||
|
||||
// Apply options to request
|
||||
if opts != nil {
|
||||
if req.Options == nil {
|
||||
req.Options = make(map[string]interface{})
|
||||
}
|
||||
|
||||
if opts.Temperature != nil {
|
||||
req.Options["temperature"] = *opts.Temperature
|
||||
}
|
||||
if opts.TopP != nil {
|
||||
req.Options["top_p"] = *opts.TopP
|
||||
}
|
||||
if opts.TopK != nil {
|
||||
req.Options["top_k"] = *opts.TopK
|
||||
}
|
||||
if opts.RepeatPenalty != nil {
|
||||
req.Options["repeat_penalty"] = *opts.RepeatPenalty
|
||||
}
|
||||
if opts.Seed != nil {
|
||||
req.Options["seed"] = *opts.Seed
|
||||
}
|
||||
if opts.Stop != nil {
|
||||
req.Options["stop"] = opts.Stop
|
||||
}
|
||||
if opts.Format != "" {
|
||||
req.Format = opts.Format
|
||||
}
|
||||
|
||||
// Merge additional options
|
||||
for k, v := range opts.Options {
|
||||
req.Options[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
return s.Generate(ctx, req)
|
||||
}
|
||||
|
||||
// Chat performs a conversational chat with the model
|
||||
func (s *Service) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Starting chat", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"message_count": len(req.Messages),
|
||||
})
|
||||
}
|
||||
|
||||
// Set defaults
|
||||
if req.Model == "" {
|
||||
req.Model = s.config.DefaultModel
|
||||
}
|
||||
|
||||
var resp ChatResponse
|
||||
err := s.client.Post(ctx, "/api/chat", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Chat failed", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, ErrChatFailed(err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Chat completed", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"response_length": len(resp.Message.Content),
|
||||
"eval_count": resp.EvalCount,
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Embed generates embeddings for the given text
|
||||
func (s *Service) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Generating embeddings", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"text_length": len(req.Prompt),
|
||||
})
|
||||
}
|
||||
|
||||
var resp EmbedResponse
|
||||
err := s.client.Post(ctx, "/api/embeddings", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Embedding generation failed", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, ErrEmbeddingFailed(err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Embeddings generated", map[string]interface{}{
|
||||
"model": req.Model,
|
||||
"embedding_dimension": len(resp.Embedding),
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// PullModel pulls a model from the registry
|
||||
func (s *Service) PullModel(ctx context.Context, model string) (*PullResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Pulling model", map[string]interface{}{
|
||||
"model": model,
|
||||
})
|
||||
}
|
||||
|
||||
req := &PullRequest{
|
||||
Name: model,
|
||||
}
|
||||
|
||||
var resp PullResponse
|
||||
err := s.client.Post(ctx, "/api/pull", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Model pull failed", map[string]interface{}{
|
||||
"model": model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, ErrModelPullFailed(model, err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Model pulled successfully", map[string]interface{}{
|
||||
"model": model,
|
||||
"status": resp.Status,
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// ListModels lists available models
|
||||
func (s *Service) ListModels(ctx context.Context) (*ListModelsResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Listing models", map[string]interface{}{})
|
||||
}
|
||||
|
||||
var resp ListModelsResponse
|
||||
err := s.client.Get(ctx, "/api/tags", &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Failed to list models", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, fmt.Errorf("failed to list models: %w", err)
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Models listed", map[string]interface{}{
|
||||
"model_count": len(resp.Models),
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// DeleteModel deletes a model
|
||||
func (s *Service) DeleteModel(ctx context.Context, model string) (*DeleteResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Deleting model", map[string]interface{}{
|
||||
"model": model,
|
||||
})
|
||||
}
|
||||
|
||||
req := &DeleteRequest{
|
||||
Name: model,
|
||||
}
|
||||
|
||||
var resp DeleteResponse
|
||||
err := s.client.Delete(ctx, "/api/delete", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Model deletion failed", map[string]interface{}{
|
||||
"model": model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, ErrModelDeleteFailed(model, err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Info("Model deleted successfully", map[string]interface{}{
|
||||
"model": model,
|
||||
"status": resp.Status,
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// ShowModel shows model information
|
||||
func (s *Service) ShowModel(ctx context.Context, model string) (*ShowModelResponse, error) {
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Showing model information", map[string]interface{}{
|
||||
"model": model,
|
||||
})
|
||||
}
|
||||
|
||||
req := &ShowModelRequest{
|
||||
Name: model,
|
||||
}
|
||||
|
||||
var resp ShowModelResponse
|
||||
err := s.client.Post(ctx, "/api/show", req, &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Failed to show model", map[string]interface{}{
|
||||
"model": model,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil, fmt.Errorf("failed to show model %s: %w", model, err)
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Model information retrieved", map[string]interface{}{
|
||||
"model": model,
|
||||
})
|
||||
}
|
||||
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Health checks if the Ollama service is available
|
||||
func (s *Service) Health(ctx context.Context) error {
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Checking Ollama health", map[string]interface{}{})
|
||||
}
|
||||
|
||||
var resp HealthResponse
|
||||
err := s.client.Get(ctx, "/api/version", &resp)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Error("Health check failed", map[string]interface{}{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return ErrConnectionFailed(err.Error())
|
||||
}
|
||||
|
||||
if s.logger != nil {
|
||||
s.logger.Debug("Health check passed", map[string]interface{}{})
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewChatBuilder creates a new chat builder for conversational interactions
|
||||
func (s *Service) NewChatBuilder() ChatBuilder {
|
||||
return &chatBuilder{
|
||||
service: s,
|
||||
request: &ChatRequest{
|
||||
Model: s.config.DefaultModel,
|
||||
Messages: make([]Message, 0),
|
||||
Options: ChatOptions{
|
||||
Temperature: &s.config.DefaultTemperature,
|
||||
TopP: &s.config.DefaultTopP,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// chatBuilder implements the ChatBuilder interface
|
||||
type chatBuilder struct {
|
||||
service OllamaService
|
||||
request *ChatRequest
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetModel(model string) ChatBuilder {
|
||||
b.request.Model = model
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetSystem(system string) ChatBuilder {
|
||||
// Add or update system message
|
||||
for i, msg := range b.request.Messages {
|
||||
if msg.Role == MessageRoleSystem {
|
||||
b.request.Messages[i].Content = system
|
||||
return b
|
||||
}
|
||||
}
|
||||
|
||||
// Add new system message at the beginning
|
||||
systemMsg := Message{
|
||||
Role: MessageRoleSystem,
|
||||
Content: system,
|
||||
}
|
||||
b.request.Messages = append([]Message{systemMsg}, b.request.Messages...)
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) AddMessage(role MessageRole, content string) ChatBuilder {
|
||||
msg := Message{
|
||||
Role: role,
|
||||
Content: content,
|
||||
}
|
||||
b.request.Messages = append(b.request.Messages, msg)
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetTemperature(temperature float64) ChatBuilder {
|
||||
b.request.Options.Temperature = &temperature
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetTopP(topP float64) ChatBuilder {
|
||||
b.request.Options.TopP = &topP
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetMaxTokens(maxTokens int) ChatBuilder {
|
||||
b.request.Options.NumPredict = &maxTokens
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetStream(stream bool) ChatBuilder {
|
||||
b.request.Stream = &stream
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetFormat(format string) ChatBuilder {
|
||||
b.request.Format = format
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) SetOptions(options map[string]interface{}) ChatBuilder {
|
||||
if b.request.Options.Options == nil {
|
||||
b.request.Options.Options = make(map[string]interface{})
|
||||
}
|
||||
for k, v := range options {
|
||||
b.request.Options.Options[k] = v
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *chatBuilder) Build() (*ChatRequest, error) {
|
||||
if len(b.request.Messages) == 0 {
|
||||
return nil, ErrInvalidRequest("at least one message is required")
|
||||
}
|
||||
return b.request, nil
|
||||
}
|
||||
|
||||
func (b *chatBuilder) Send(ctx context.Context) (*ChatResponse, error) {
|
||||
req, err := b.Build()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b.service.Chat(ctx, req)
|
||||
}
|
||||
|
||||
func (b *chatBuilder) Continue(ctx context.Context, userMessage string) (*ChatResponse, error) {
|
||||
// Add user message
|
||||
b.AddMessage(MessageRoleUser, userMessage)
|
||||
|
||||
// Send chat request
|
||||
resp, err := b.Send(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Add assistant response to conversation history
|
||||
b.AddMessage(MessageRoleAssistant, resp.Message.Content)
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
@@ -2,7 +2,7 @@ package schedule
|
||||
|
||||
import (
|
||||
"time"
|
||||
"tm/internal/tender"
|
||||
"tm/internal/notice"
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
type CronScheduler struct {
|
||||
cron *cron.Cron
|
||||
mongoManager *mongo.ConnectionManager
|
||||
tenderRepo tender.TenderRepository
|
||||
noticeRepo notice.Repository
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
@@ -31,14 +31,14 @@ type Job struct {
|
||||
}
|
||||
|
||||
// NewCronScheduler creates a new cron scheduler instance
|
||||
func NewCronScheduler(logger logger.Logger, mongoManager *mongo.ConnectionManager, tenderRepo tender.TenderRepository) *CronScheduler {
|
||||
func NewCronScheduler(logger logger.Logger, mongoManager *mongo.ConnectionManager, noticeRepo notice.Repository) *CronScheduler {
|
||||
// Create cron with timezone support and seconds precision
|
||||
c := cron.New(cron.WithLocation(time.UTC), cron.WithSeconds())
|
||||
|
||||
return &CronScheduler{
|
||||
cron: c,
|
||||
mongoManager: mongoManager,
|
||||
tenderRepo: tenderRepo,
|
||||
noticeRepo: noticeRepo,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
+21
-21
@@ -4,11 +4,11 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"tm/internal/tender"
|
||||
"tm/internal/notice"
|
||||
)
|
||||
|
||||
// mapToTenderFromParsedDoc maps a parsed TED document to tender entity
|
||||
func (s *TEDScraper) mapToTenderFromParsedDoc(parsedDoc *ParsedDocument, sourceFileURL, sourceFileName string) *tender.Tender {
|
||||
func (s *TEDScraper) mapToTenderFromParsedDoc(parsedDoc *ParsedDocument, sourceFileURL, sourceFileName string) *notice.Notice {
|
||||
switch parsedDoc.Type {
|
||||
case DocumentTypeContractNotice:
|
||||
if parsedDoc.ContractNotice == nil {
|
||||
@@ -45,7 +45,7 @@ func (s *TEDScraper) mapToTenderFromParsedDoc(parsedDoc *ParsedDocument, sourceF
|
||||
}
|
||||
|
||||
// mapGenericNoticeToTender maps any TED document type to tender entity using the common interface
|
||||
func (s *TEDScraper) mapGenericNoticeToTender(doc TEDDocument, sourceFileURL, sourceFileName string) *tender.Tender {
|
||||
func (s *TEDScraper) mapGenericNoticeToTender(doc TEDDocument, sourceFileURL, sourceFileName string) *notice.Notice {
|
||||
if doc == nil {
|
||||
s.logger.Error("Document is nil in mapGenericNoticeToTender", map[string]interface{}{})
|
||||
return nil
|
||||
@@ -53,7 +53,7 @@ func (s *TEDScraper) mapGenericNoticeToTender(doc TEDDocument, sourceFileURL, so
|
||||
|
||||
now := time.Now().Unix()
|
||||
|
||||
t := &tender.Tender{
|
||||
t := ¬ice.Notice{
|
||||
ContractNoticeID: doc.GetID(),
|
||||
ContractFolderID: "", // Not available in generic interface
|
||||
NoticeTypeCode: doc.GetNoticeType(),
|
||||
@@ -68,11 +68,11 @@ func (s *TEDScraper) mapGenericNoticeToTender(doc TEDDocument, sourceFileURL, so
|
||||
// Basic tender information
|
||||
Title: "Generic Notice " + doc.GetNoticeType(),
|
||||
Description: "Processed from " + doc.GetNoticeType() + " document",
|
||||
Status: tender.TenderStatusActive,
|
||||
Source: tender.TenderSourceTEDScraper,
|
||||
Status: notice.TenderStatusActive,
|
||||
Source: notice.TenderSourceTEDScraper,
|
||||
SourceFileURL: sourceFileURL,
|
||||
SourceFileName: sourceFileName,
|
||||
ProcessingMetadata: tender.ProcessingMetadata{
|
||||
ProcessingMetadata: notice.ProcessingMetadata{
|
||||
ScrapedAt: now,
|
||||
ProcessedAt: now,
|
||||
ProcessingVersion: "1.0",
|
||||
@@ -86,10 +86,10 @@ func (s *TEDScraper) mapGenericNoticeToTender(doc TEDDocument, sourceFileURL, so
|
||||
}
|
||||
|
||||
// mapToTender maps a TED ContractNotice to tender entity
|
||||
func (s *TEDScraper) mapToTender(cn *ContractNotice, sourceFileURL, sourceFileName string) *tender.Tender {
|
||||
func (s *TEDScraper) mapToTender(cn *ContractNotice, sourceFileURL, sourceFileName string) *notice.Notice {
|
||||
now := time.Now().Unix()
|
||||
|
||||
t := &tender.Tender{
|
||||
t := ¬ice.Notice{
|
||||
ContractNoticeID: cn.ID,
|
||||
ContractFolderID: cn.ContractFolderID,
|
||||
NoticeTypeCode: cn.GetNoticeSubType(),
|
||||
@@ -105,11 +105,11 @@ func (s *TEDScraper) mapToTender(cn *ContractNotice, sourceFileURL, sourceFileNa
|
||||
MainClassification: cn.GetMainClassification(),
|
||||
PlaceOfPerformance: cn.GetPlaceOfPerformance(),
|
||||
DocumentURI: cn.GetDocumentURI(),
|
||||
Status: tender.TenderStatusActive,
|
||||
Source: tender.TenderSourceTEDScraper,
|
||||
Status: notice.TenderStatusActive,
|
||||
Source: notice.TenderSourceTEDScraper,
|
||||
SourceFileURL: sourceFileURL,
|
||||
SourceFileName: sourceFileName,
|
||||
ProcessingMetadata: tender.ProcessingMetadata{
|
||||
ProcessingMetadata: notice.ProcessingMetadata{
|
||||
ScrapedAt: now,
|
||||
ProcessedAt: now,
|
||||
ProcessingVersion: "1.0",
|
||||
@@ -215,7 +215,7 @@ func (s *TEDScraper) mapToTender(cn *ContractNotice, sourceFileURL, sourceFileNa
|
||||
// Set selection criteria
|
||||
if criteria := cn.GetSelectionCriteria(); criteria != nil {
|
||||
for _, criterion := range criteria {
|
||||
t.SelectionCriteria = append(t.SelectionCriteria, tender.SelectionCriterion{
|
||||
t.SelectionCriteria = append(t.SelectionCriteria, notice.SelectionCriterion{
|
||||
TypeCode: criterion.CriterionTypeCode,
|
||||
Description: criterion.Description,
|
||||
LanguageID: criterion.LanguageID,
|
||||
@@ -235,10 +235,10 @@ func (s *TEDScraper) mapToTender(cn *ContractNotice, sourceFileURL, sourceFileNa
|
||||
}
|
||||
|
||||
// mapContractAwardNoticeToTender maps a TED ContractAwardNotice to tender entity
|
||||
func (s *TEDScraper) mapContractAwardNoticeToTender(can *ContractAwardNotice, sourceFileURL, sourceFileName string) *tender.Tender {
|
||||
func (s *TEDScraper) mapContractAwardNoticeToTender(can *ContractAwardNotice, sourceFileURL, sourceFileName string) *notice.Notice {
|
||||
now := time.Now().Unix()
|
||||
|
||||
t := &tender.Tender{
|
||||
t := ¬ice.Notice{
|
||||
ContractNoticeID: can.ID,
|
||||
ContractFolderID: can.ContractFolderID,
|
||||
NoticeTypeCode: can.NoticeTypeCode,
|
||||
@@ -252,11 +252,11 @@ func (s *TEDScraper) mapContractAwardNoticeToTender(can *ContractAwardNotice, so
|
||||
ProcurementTypeCode: can.GetContractNature(),
|
||||
ProcedureCode: can.GetProcedureCode(),
|
||||
MainClassification: can.GetMainClassification(),
|
||||
Status: tender.TenderStatusAwarded, // Award notices are for completed tenders
|
||||
Source: tender.TenderSourceTEDScraper,
|
||||
Status: notice.TenderStatusAwarded, // Award notices are for completed tenders
|
||||
Source: notice.TenderSourceTEDScraper,
|
||||
SourceFileURL: sourceFileURL,
|
||||
SourceFileName: sourceFileName,
|
||||
ProcessingMetadata: tender.ProcessingMetadata{
|
||||
ProcessingMetadata: notice.ProcessingMetadata{
|
||||
ScrapedAt: now,
|
||||
ProcessedAt: now,
|
||||
ProcessingVersion: "1.0",
|
||||
@@ -315,13 +315,13 @@ func (s *TEDScraper) mapContractAwardNoticeToTender(can *ContractAwardNotice, so
|
||||
}
|
||||
|
||||
// mapOrganization maps a TED Organization to tender Organization
|
||||
func (s *TEDScraper) mapOrganization(tedOrg *Organization, role string) *tender.Organization {
|
||||
func (s *TEDScraper) mapOrganization(tedOrg *Organization, role string) *notice.Organization {
|
||||
if tedOrg == nil || tedOrg.Company == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
company := tedOrg.Company
|
||||
org := &tender.Organization{
|
||||
org := ¬ice.Organization{
|
||||
Role: role,
|
||||
}
|
||||
|
||||
@@ -351,7 +351,7 @@ func (s *TEDScraper) mapOrganization(tedOrg *Organization, role string) *tender.
|
||||
|
||||
// Set address
|
||||
if company.PostalAddress != nil {
|
||||
org.Address = tender.Address{
|
||||
org.Address = notice.Address{
|
||||
StreetName: company.PostalAddress.StreetName,
|
||||
CityName: company.PostalAddress.CityName,
|
||||
PostalZone: company.PostalAddress.PostalZone,
|
||||
|
||||
+39
-35
@@ -10,7 +10,7 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
"tm/internal/tender"
|
||||
"tm/internal/notice"
|
||||
"tm/pkg/logger"
|
||||
"tm/pkg/mongo"
|
||||
"tm/pkg/notification"
|
||||
@@ -27,6 +27,7 @@ type Config struct {
|
||||
RetryDelay time.Duration `mapstructure:"retry_delay"`
|
||||
CleanupAfter time.Duration `mapstructure:"cleanup_after"`
|
||||
ScrapingInterval string `mapstructure:"scraping_interval"`
|
||||
AlertMail string `mapstructure:"ALERT_MAIL"`
|
||||
}
|
||||
|
||||
// TEDScraper handles downloading and parsing TED XML files
|
||||
@@ -37,7 +38,7 @@ type TEDScraper struct {
|
||||
httpClient *http.Client
|
||||
xmlParser *TEDParser
|
||||
mongoManager *mongo.ConnectionManager
|
||||
tenderRepo tender.TenderRepository
|
||||
tenderRepo notice.Repository
|
||||
}
|
||||
|
||||
// NewTEDScraper creates a new TED scraper instance
|
||||
@@ -45,7 +46,7 @@ func NewTEDScraper(
|
||||
config *Config,
|
||||
logger logger.Logger,
|
||||
mongoManager *mongo.ConnectionManager,
|
||||
tenderRepo tender.TenderRepository,
|
||||
tenderRepo notice.Repository,
|
||||
notify notification.SDK,
|
||||
) *TEDScraper {
|
||||
httpClient := &http.Client{
|
||||
@@ -251,6 +252,8 @@ func (s *TEDScraper) tarGzFileProcessor(ctx context.Context, reader io.Reader, o
|
||||
}
|
||||
}
|
||||
|
||||
if s.config.AlertMail != "" {
|
||||
|
||||
go s.notify.SendNotification(
|
||||
context.Background(),
|
||||
¬ification.NotificationRequest{
|
||||
@@ -266,16 +269,17 @@ func (s *TEDScraper) tarGzFileProcessor(ctx context.Context, reader io.Reader, o
|
||||
}),
|
||||
Priority: "important",
|
||||
Methods: notification.NotificationMethods{
|
||||
Email: "nakhostin.nima1998@gmail.com",
|
||||
Email: s.config.AlertMail,
|
||||
},
|
||||
UserID: "ted-scraper",
|
||||
Type: "alert",
|
||||
})
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *TEDScraper) findTenderByContractNoticeID(ctx context.Context, contractNoticeID string) (*tender.Tender, error) {
|
||||
func (s *TEDScraper) findTenderByContractNoticeID(ctx context.Context, contractNoticeID string) (*notice.Notice, error) {
|
||||
existingTender, err := s.tenderRepo.GetByContractNoticeID(ctx, contractNoticeID)
|
||||
if err != nil {
|
||||
if err.Error() == "document not found" {
|
||||
@@ -308,37 +312,37 @@ func (s *TEDScraper) processXMLContent(ctx context.Context, content []byte, file
|
||||
// Map to tender entity
|
||||
t := s.mapToTenderFromParsedDoc(notice, "", fileName)
|
||||
|
||||
existingTender, err := s.findTenderByContractNoticeID(ctx, t.ContractNoticeID)
|
||||
if err != nil {
|
||||
s.logger.Warn("Error checking for existing tender", map[string]interface{}{
|
||||
"contract_notice_id": t.ContractNoticeID,
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
// existingTender, err := s.findTenderByContractNoticeID(ctx, t.ContractNoticeID)
|
||||
// if err != nil {
|
||||
// s.logger.Warn("Error checking for existing tender", map[string]interface{}{
|
||||
// "contract_notice_id": t.ContractNoticeID,
|
||||
// "error": err.Error(),
|
||||
// })
|
||||
// }
|
||||
|
||||
if err == nil && existingTender != nil {
|
||||
// Update existing tender
|
||||
s.logger.Info("Found existing tender, updating", map[string]interface{}{
|
||||
"existing_tender_id": existingTender.ID.Hex(),
|
||||
"contract_notice_id": t.ContractNoticeID,
|
||||
})
|
||||
// if err == nil && existingTender != nil {
|
||||
// // Update existing tender
|
||||
// s.logger.Info("Found existing tender, updating", map[string]interface{}{
|
||||
// "existing_tender_id": existingTender.ID.Hex(),
|
||||
// "contract_notice_id": t.ContractNoticeID,
|
||||
// })
|
||||
|
||||
t.ID = existingTender.ID
|
||||
t.UpdatedAt = time.Now().Unix()
|
||||
err = s.tenderRepo.Update(ctx, t)
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to update existing tender", map[string]interface{}{
|
||||
"tender_id": t.ID.Hex(),
|
||||
"error": err.Error(),
|
||||
})
|
||||
return fmt.Errorf("failed to update existing tender: %w", err)
|
||||
}
|
||||
// t.ID = existingTender.ID
|
||||
// t.UpdatedAt = time.Now().Unix()
|
||||
// err = s.tenderRepo.Update(ctx, t)
|
||||
// if err != nil {
|
||||
// s.logger.Error("Failed to update existing tender", map[string]interface{}{
|
||||
// "tender_id": t.ID.Hex(),
|
||||
// "error": err.Error(),
|
||||
// })
|
||||
// return fmt.Errorf("failed to update existing tender: %w", err)
|
||||
// }
|
||||
|
||||
s.logger.Info("Successfully updated tender", map[string]interface{}{
|
||||
"tender_id": t.ID.Hex(),
|
||||
})
|
||||
} else {
|
||||
// Create new tender
|
||||
// s.logger.Info("Successfully updated tender", map[string]interface{}{
|
||||
// "tender_id": t.ID.Hex(),
|
||||
// })
|
||||
// } else {
|
||||
// // Create new tender
|
||||
s.logger.Info("Creating new tender", map[string]interface{}{
|
||||
"contract_notice_id": t.ContractNoticeID,
|
||||
"tender_id": t.TenderID,
|
||||
@@ -347,7 +351,7 @@ func (s *TEDScraper) processXMLContent(ctx context.Context, content []byte, file
|
||||
t.CreatedAt = time.Now().Unix()
|
||||
t.UpdatedAt = t.CreatedAt
|
||||
|
||||
err = s.tenderRepo.Create(ctx, t)
|
||||
err = s.tenderRepo.Import(ctx, t)
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to create new tender", map[string]interface{}{
|
||||
"contract_notice_id": t.ContractNoticeID,
|
||||
@@ -361,7 +365,7 @@ func (s *TEDScraper) processXMLContent(ctx context.Context, content []byte, file
|
||||
"tender_id": t.TenderID,
|
||||
"contract_notice_id": t.ContractNoticeID,
|
||||
})
|
||||
}
|
||||
// }
|
||||
|
||||
// Log detailed parsed information based on notice type
|
||||
s.logger.Info("Successfully parsed notice", map[string]interface{}{
|
||||
|
||||
+2
-2
@@ -5,7 +5,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"tm/internal/tender"
|
||||
"tm/internal/notice"
|
||||
)
|
||||
|
||||
// parseDate parses various date formats used in TED XML
|
||||
@@ -146,7 +146,7 @@ func GenerateTenderURL(noticePublicationID string) string {
|
||||
|
||||
// generateTenderID generates a unique tender ID using the format:
|
||||
// {Source}{BUYER_CODE}{TED_CODE}{DATE}{LOCATION_CODE}
|
||||
func GenerateTenderID(t *tender.Tender) string {
|
||||
func GenerateTenderID(t *notice.Notice) string {
|
||||
var parts []string
|
||||
|
||||
// Source (TED)
|
||||
|
||||
Reference in New Issue
Block a user